diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 184bf0de275..f1fa762d772 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -189,6 +189,16 @@ Spark's monitor and would wait for the consumer's, while whatever held the consu Spark's. For the same reason, nothing reachable from a `spill` callback may allocate memory that routes back through the same consumer. +**A short grant must not take the `TaskMemoryManager` monitor again.** `acquireExecutionMemory` +holds that monitor for the whole call, including while `ExecutionMemoryPool.acquireMemory` waits in +`lock.wait()`, which gives up the memory manager's monitor but not the task's. When +`CometTaskMemoryManager.acquireMemory` gets less than it asked for, its thread holds those bytes +until native code hands them back, and another acquire of the same task can be waiting for exactly +those bytes. If the first thread then took the task's monitor, for example through +`TaskMemoryManager.showMemoryUsage`, neither thread could go on until an unrelated task freed +memory. So a short grant is logged with the figures already in hand and +`getMemoryConsumptionForThisTask`, which takes only the memory manager's monitor. + **`scala.util.control.NonFatal` does not contain an acquisition.** `acquireExecutionMemory` fails in three ways and `NonFatal` catches only the first. It runs other consumers' `spill`, where `TaskMemoryManager` turns an interrupted spill into a `RuntimeException` and an `IOException` into a @@ -276,11 +286,12 @@ ignored and the pool is always `UnboundedMemoryPool`. the inside out, a Comet plan in the default configuration sees: ```text -[LoggingMemoryPool] <- only when spark.comet.debug.memory=true - [TaskSharedMemoryPool] <- RAII handle for the per-task registry - [SpillReplayPool] <- lets a spilled final aggregate read its spill files back - [TrackConsumersPool] <- DataFusion; names the top 10 consumers in error messages - [CometFairMemoryPool] <- delegates acquire/release to Spark over JNI +[PlanMemoryPool] <- counts the plan's bytes so releasing the plan can wait for them + [LoggingMemoryPool] <- only when spark.comet.debug.memory=true + [TaskSharedMemoryPool] <- RAII handle for the per-task registry + [SpillReplayPool] <- lets a spilled final aggregate read its spill files back + [TrackConsumersPool] <- DataFusion; names the top 10 consumers in error messages + [CometFairMemoryPool] <- delegates acquire/release to Spark over JNI ``` Each decorator forwards every `MemoryPool` method to its inner pool, so `reserved()` at any level @@ -307,8 +318,14 @@ JNI, which goes through Spark's ordinary `TaskMemoryManager`. That means: - `grow` cannot fail, so when Spark grants less than asked the pool records the full amount anyway and carries the shortfall as _overcommit_ (`spark_memory.rs`). A later `shrink` repays the overcommit before returning anything to Spark, so Spark is never handed back more than it - granted. While overcommit is outstanding, `try_grow` asks Spark for it on top of the request and - is refused unless Spark can cover both, so operators spill until the debt is repaid. Both pools' + granted. Each `try_grow` asks Spark for the debt no other `try_grow` in flight is already asking + for, on top of its own request, and is refused unless Spark covers both. A release that repays + the debt while such a request is parked in Spark leaves the request larger than needed, because + Spark cannot shrink a parked request. It waits until other memory frees or the number of active + tasks changes, so the task can be held up until an unrelated task releases memory, and the + request may then come back short and be refused. Any request that carries debt has this window, + with or without a second `try_grow`. Closing it would mean handing claimed debt to Spark on + release, which gives Spark back bytes it never granted if the claim is then refused. Both pools' `Display` output and their `try_grow` errors report the current overcommit. `CometFairMemoryPool` additionally applies two local checks before it asks Spark, and refuses the @@ -340,6 +357,41 @@ This is why `fair_unified` can spill earlier than `greedy_unified`: a consumer a refused even when the rest of the pool is free, which keeps that memory for the task's other consumers. +**The pool mutex is never held across a JNI call.** Both checks run, and the bytes are charged to +the pool's total and to the consumer's running total, under the lock. The lock is dropped before +`acquireMemory` or `releaseMemory` runs, and the bookkeeping is settled after the call returns. This +holds for `grow` as well as `try_grow`: `grow` skips both checks but charges its bytes under the +lock the same way, and its JVM call runs without it. The reason is an acquire parked inside +`ExecutionMemoryPool.acquireMemory` below its minimum share: it waits for another thread of the same +task to release memory. If the release had to take a lock that the parked acquire was holding, it +could not land until an unrelated task freed memory, so the task would stall meanwhile. Such a +release can empty the task's balance, which makes Spark drop the task's entry under the parked +acquire, and `CometTaskMemoryManager.acquireMemory` then asks Spark again. + +Settling after the call leaves two windows, both on the conservative side: + +- A `try_grow` is charged to the pool's total and its consumer's when it passes both checks, and a + `grow` when it is called, before Spark answers. A second `try_grow` in that window is checked + against totals that include the first, so it can be refused where waiting for the first to settle + would have let it through. This charge is what keeps two concurrent grows from jointly taking a + consumer past its share or the pool past `pool_size`. +- A shrink, and the rollback of a short grant, hand the bytes to Spark first and take them off the + pool's total and the consumer's only once Spark has them. A grow in that window is checked against + totals that still include those bytes, so a grow that only fits once they are free is refused by + the share or pool check instead of being sent to Spark ahead of the release. A grow that fits + without them still goes to Spark and can come back with a short grant if the release has not + landed and the task is at its Spark share. If the rollback release itself fails, those bytes stay + charged for the pool's life, because Spark holds them until the task ends. + +Neither window admits a `try_grow` that the two checks would have refused. A refusal is the ordinary +`ResourcesExhausted` error, which a spillable operator answers by spilling. `grow` is never refused: +it is not subject to either check, and whatever Spark declines of it becomes overcommit, as +described above. + +The overcommit window in [The unified pools](#the-unified-pools) never admits a request these checks +would refuse either, but unlike these two it can delay one: it can hold a `try_grow` until an +unrelated task frees memory. It applies to both pools. + ### Task-shared pools and their lifetime A single Spark task can run more than one native plan at a time. A native shuffle is not one of diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 262d8b8fd41..117c064e833 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -171,11 +171,11 @@ fn log_native_allocated() { /// Used to sum memory reservations across all contexts for the memory usage log and tracing. /// /// Never read a pool's reservation while holding this registry's lock; copy the pools out with -/// [`snapshot_registry`] and read them after it is released. `CometFairMemoryPool` holds its own -/// lock across the JNI call that acquires memory from Spark, and Spark can park that call until -/// another task frees memory. A finishing task frees its reservations only after `releasePlan` has -/// taken this lock to unregister, so a reservation read under this lock can wait on a pool that is -/// itself waiting on the lock. +/// [`snapshot_registry`] and read them after it is released. `reserved()` takes the pool's own +/// lock, and this registry must never wait on a pool: `releasePlan` takes the registry lock to +/// unregister while the task's other plans keep reserving and releasing, so a pool lock taken +/// under it would put every context on the executor behind that one pool. +/// `reservations_are_read_outside_the_registry_lock` pins the rule. type ThreadPoolMap = HashMap>>; static THREAD_MEMORY_POOLS: OnceLock> = OnceLock::new(); @@ -2676,8 +2676,8 @@ mod tests { assert_eq!(memory_usage().pools_reserved - before.pools_reserved, 3072); } - /// Stands in for a `CometFairMemoryPool` whose lock is held across a Spark acquire: it counts - /// its reservation reads, and notes whether the registry lock was held during any of them. + /// Stands in for a pool whose `reserved()` takes its own lock: it counts its reservation + /// reads, and notes whether the registry lock was held during any of them. #[derive(Debug, Default)] struct RegistryProbePool { reads: std::sync::atomic::AtomicUsize, diff --git a/native/core/src/execution/memory_pools/fair_pool.rs b/native/core/src/execution/memory_pools/fair_pool.rs index 56b09da7365..f46eba3a200 100644 --- a/native/core/src/execution/memory_pools/fair_pool.rs +++ b/native/core/src/execution/memory_pools/fair_pool.rs @@ -20,6 +20,8 @@ use std::{ fmt::{Debug, Display, Formatter, Result as FmtResult}, }; +use log::warn; + use super::spark_memory::SparkMemory; use datafusion::common::resources_err; use datafusion::execution::memory_pool::MemoryConsumer; @@ -31,6 +33,11 @@ use parking_lot::Mutex; /// A DataFusion fair `MemoryPool` implementation for Comet. Internally this is /// implemented via delegating calls to [`crate::jvm_bridge::CometTaskMemoryManager`]. +/// +/// No lock is held across a JVM call, so a release can land while an acquire of the same task +/// waits inside Spark. If that release takes the task's balance to zero, Spark drops the task's +/// `memoryForTask` entry and the waiting acquire fails when it wakes; `CometTaskMemoryManager` +/// asks Spark again in that case. pub struct CometFairMemoryPool { spark: SparkMemory, pool_size: usize, @@ -43,17 +50,38 @@ struct CometFairPoolState { /// reservations that `new_empty()`, `split()` and `take()` create belong to the same consumer, /// so they count against one share. The pool keeps these totals itself because /// `reservation.size()` covers only one reservation, and DataFusion updates it after calling - /// `try_grow` but before calling `shrink`. + /// `try_grow` but before calling `shrink`. Bytes are charged here and to `used` together, and + /// settle together once the JVM has answered. consumers: HashMap, } impl CometFairPoolState { - /// The bytes held by `reservation`'s consumer across all of its reservations. - fn consumer_used(&mut self, reservation: &MemoryReservation) -> &mut usize { + /// The bytes held by `consumer` across all of its reservations. + fn consumer_used(&mut self, consumer: usize) -> &mut usize { self.consumers - .get_mut(&reservation.consumer().id()) + .get_mut(&consumer) .expect("reservation's consumer is not registered with the pool") } + + /// Charges `bytes` to the pool's total and to `consumer`'s share. + fn charge(&mut self, consumer: usize, bytes: usize) { + self.used = self.used.saturating_add(bytes); + let consumer_used = self.consumer_used(consumer); + *consumer_used = consumer_used.saturating_add(bytes); + } + + /// Takes `charged` bytes off the pool's total and `consumer`'s share and puts `held` back. + fn settle(&mut self, consumer: usize, charged: usize, held: usize) { + let settle = |total: usize| { + total + .checked_sub(charged) + .and_then(|total| total.checked_add(held)) + .expect("settled more bytes than the pool tracks") + }; + self.used = settle(self.used); + let consumer_used = self.consumer_used(consumer); + *consumer_used = settle(*consumer_used); + } } impl Debug for CometFairMemoryPool { @@ -84,6 +112,19 @@ impl CometFairMemoryPool { pub(super) fn overcommit(&self) -> usize { self.spark.overcommit() } + + /// Settles a release the JVM has accepted: the bytes come off the pool's total and the + /// consumer's share only now, so a grow is never admitted on bytes Spark still holds. + fn settle_release(&self, consumer: usize, bytes: usize) { + self.state.lock().settle(consumer, bytes, 0); + } + + /// Settles a finished JVM acquire. `charged` bytes went on the pool's total and the + /// consumer's share before the call and `held` is what Spark granted and still holds, which + /// stays charged until it is handed back. The difference is rolled back. + fn settle_acquire(&self, consumer: usize, charged: usize, held: usize) { + self.state.lock().settle(consumer, charged, held); + } } impl Display for CometFairMemoryPool { @@ -116,33 +157,51 @@ impl MemoryPool for CometFairMemoryPool { } /// Records memory that already exists, so it must not fail and ignores the fair and pool - /// limits. - /// See [`SparkMemory`]. + /// limits. What Spark declines is carried as overcommit, see [`SparkMemory`]. fn grow(&self, reservation: &MemoryReservation, additional: usize) { if additional == 0 { return; } - let mut state = self.state.lock(); - self.spark.acquire(additional); - state.used = state.used.saturating_add(additional); - let consumer_used = state.consumer_used(reservation); - *consumer_used = consumer_used.saturating_add(additional); + let consumer = reservation.consumer().id(); + // Charged before the JVM call, as in try_grow, so a try_grow racing this one is checked + // against totals that already include these bytes. The JVM call then runs without the + // lock held. + self.state.lock().charge(consumer, additional); + let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + self.spark.acquire(additional); + })); + if let Err(panic) = outcome { + // The caller's reservation never records bytes its grow panicked on, so the charge + // must not outlive the panic either. + self.settle_acquire(consumer, additional, 0); + std::panic::resume_unwind(panic); + } } fn shrink(&self, reservation: &MemoryReservation, subtractive: usize) { if subtractive > 0 { - let mut state = self.state.lock(); - let consumer_used = *state.consumer_used(reservation); - if consumer_used < subtractive { - panic!( - "Failed to release {subtractive} bytes where only {consumer_used} bytes tracked for the consumer" - ) + let consumer = reservation.consumer().id(); + { + let consumer_used = *self.state.lock().consumer_used(consumer); + if consumer_used < subtractive { + panic!( + "Failed to release {subtractive} bytes where only {consumer_used} bytes tracked for the consumer" + ) + } } + // The JVM release runs without the lock so a blocked acquire on another thread can + // never stall this release. The bytes stay charged until Spark has them back, so a + // grow racing this shrink is not admitted on them and sent to Spark ahead of the + // release. Outstanding overcommit is repaid before anything goes back to Spark. A + // failed release here panics (the caller already gave the bytes up, there is no one + // left to handle an error), while the short-grant path in try_grow returns Err + // because its caller can still spill. The check above is advisory once the lock is + // dropped: DataFusion never shrinks a reservation past its size, which is what keeps + // the subtraction in settle_release in range. self.spark .release(subtractive) .unwrap_or_else(|_| panic!("Failed to release {subtractive} bytes")); - state.used = state.used.checked_sub(subtractive).unwrap(); - *state.consumer_used(reservation) -= subtractive; + self.settle_release(consumer, subtractive); } } @@ -152,45 +211,86 @@ impl MemoryPool for CometFairMemoryPool { additional: usize, ) -> Result<(), DataFusionError> { if additional > 0 { - let mut state = self.state.lock(); - let num = state.consumers.len(); - let limit = self - .pool_size - .checked_div(num) - .expect("overflow in checked_div"); - let consumer_used = *state.consumer_used(reservation); - if limit < consumer_used.saturating_add(additional) { - return resources_err!( - "Failed to acquire {additional} bytes where this consumer already holds {consumer_used} bytes and the fair limit is {limit} bytes, {num} registered ({} bytes overcommitted)", - self.spark.overcommit() - ); - } - // The shares alone do not bound the pool's total, because a consumer keeps what it - // reserved before another consumer registered. - let used = state.used; - if self.pool_size < used.saturating_add(additional) { - return resources_err!( - "Failed to acquire {additional} bytes where {used} bytes already reserved ({} bytes overcommitted) and the pool limit is {} bytes", - self.spark.overcommit(), - self.pool_size - ); + let consumer = reservation.consumer().id(); + // Checking both limits and charging the bytes is one atomic step, so concurrent grows + // can never jointly take a consumer past its share or the pool past pool_size. The + // blocking JVM calls then run without any lock held, and the charge rolls back if + // the JVM does not back it. + { + let mut state = self.state.lock(); + let num = state.consumers.len(); + let limit = self + .pool_size + .checked_div(num) + .expect("overflow in checked_div"); + let consumer_used = *state.consumer_used(consumer); + if limit < consumer_used.saturating_add(additional) { + return resources_err!( + "Failed to acquire {additional} bytes where this consumer already holds {consumer_used} bytes and the fair limit is {limit} bytes, {num} registered ({} bytes overcommitted)", + self.spark.overcommit() + ); + } + // The shares alone do not bound the pool's total, because a consumer keeps what + // it reserved before another consumer registered. + let used = state.used; + if self.pool_size < used.saturating_add(additional) { + return resources_err!( + "Failed to acquire {additional} bytes where {used} bytes already reserved ({} bytes overcommitted) and the pool limit is {} bytes", + self.spark.overcommit(), + self.pool_size + ); + } + state.charge(consumer, additional); } - // A partial grant is handed back and refused, which triggers spilling in the caller. - if let Err(refusal) = self.spark.try_acquire(additional)? { - return resources_err!( - "Failed to acquire {} bytes plus {} bytes overcommitted, only got {} bytes. Reserved: {} bytes", - additional, - refusal.overcommit, - refusal.granted, - state.used - ); + // Spark is asked for the request plus the overcommit no other try_grow in flight is + // asking for, and a full grant repays it. A short grant stays with Spark until this + // pool hands it back below, so the bytes can stay charged meanwhile. The JVM call can + // panic inside its JNI frame; the optimistic reservation must not outlive it, or the + // leaked bytes poison the task-shared pool for every other consumer. + let refusal = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + self.spark.try_acquire_leaving_a_short_grant(additional) + })) { + Ok(Ok(Ok(()))) => return Ok(()), + Ok(Ok(Err(refusal))) => refusal, + Ok(Err(e)) => { + self.settle_acquire(consumer, additional, 0); + return Err(e.into()); + } + Err(panic) => { + self.settle_acquire(consumer, additional, 0); + std::panic::resume_unwind(panic); + } + }; + // A grant that falls short of the request is handed back whole and reported so the + // caller can spill. The bytes Spark did not back come off the books at once. The + // granted bytes, which can exceed the request when overcommit was asked for on top + // of it, stay charged until Spark has them back, like a shrink, so no grow is + // admitted on them meanwhile. A failed return leaves Spark holding them until the + // task ends, and they stay charged here to match. The caller still gets the short + // grant error below so a spillable operator spills. A panic in that release leaves + // the same state as an error. + let granted = refusal.granted; + self.settle_acquire(consumer, additional, granted); + if granted > 0 { + if let Err(e) = self.spark.manager().release(granted) { + warn!( + "Task {} failed to return a short grant of {granted} bytes, which stay charged to the pool: {e:?}", + self.spark.task_attempt_id() + ); + } else { + self.settle_release(consumer, granted); + } } - state.used = state - .used - .checked_add(additional) - .expect("overflow in checked_add"); - *state.consumer_used(reservation) += additional; + + return resources_err!( + "Failed to acquire {} bytes plus {} of {} bytes overcommitted, only got {} bytes. Reserved: {} bytes", + additional, + refusal.overcommit, + self.spark.overcommit(), + granted, + self.reserved() + ); } Ok(()) } @@ -202,9 +302,382 @@ impl MemoryPool for CometFairMemoryPool { #[cfg(test)] mod tests { - use super::super::spark_memory::fake::FakeSpark; + use super::super::spark_memory::{fake::FakeSpark, SparkMemoryManager}; use super::*; - use std::sync::Arc; + use crate::errors::{CometError, CometResult}; + use parking_lot::Condvar; + use std::collections::{hash_map::Entry, HashMap}; + use std::sync::atomic::{AtomicBool, AtomicI64, AtomicUsize, Ordering::SeqCst}; + use std::sync::mpsc::{channel, Receiver, Sender}; + use std::sync::{Arc, Barrier}; + use std::thread; + use std::time::Duration; + + const MIB: usize = 1 << 20; + const GIB: usize = 1 << 30; + + /// The task this pool belongs to and neighbours that compete with it for the same pool. + const THIS_TASK: i64 = 0; + const OTHER_TASK: i64 = 1; + const THIRD_TASK: i64 = 2; + + /// Real Spark parks a starved acquire forever; the stub gives up after this long and fails + /// the test instead, so a deadlock shows up as a panic rather than a hung test binary. A + /// heavily loaded CI box may need this lengthened. + const WAIT_TIMEOUT: Duration = Duration::from_secs(5); + const TEST_TIMEOUT: Duration = Duration::from_secs(20); + + /// A pause point inside a bridge call. While armed, the call announces itself on `entered` + /// and blocks until the test opens the gate, so a test can interleave a second thread at a + /// precise moment of a JNI call. + struct Gate { + armed: AtomicBool, + entered: (Sender<()>, Mutex>), + open: (Sender<()>, Mutex>), + } + + impl Gate { + fn new() -> Self { + let entered = channel(); + let open = channel(); + Self { + armed: AtomicBool::new(false), + entered: (entered.0, Mutex::new(entered.1)), + open: (open.0, Mutex::new(open.1)), + } + } + + fn arm(&self) { + self.armed.store(true, SeqCst); + } + + fn disarm(&self) { + self.armed.store(false, SeqCst); + } + + /// Called from the bridge thread. + fn pass(&self) { + if self.armed.load(SeqCst) { + let _ = self.entered.0.send(()); + // A test that fails before opening the gate must not hang the bridge thread. + let _ = self.open.1.lock().recv_timeout(TEST_TIMEOUT); + } + } + + fn wait_entered(&self, what: &str) { + self.entered + .1 + .lock() + .recv_timeout(TEST_TIMEOUT) + .unwrap_or_else(|_| panic!("{what} never reached the gate")); + } + + fn open(&self) { + let _ = self.open.0.send(()); + } + } + + /// The state Spark's `ExecutionMemoryPool` keeps under its `lock` monitor. + struct SparkPool { + pool_size: i64, + /// `memoryForTask`: created on a task's first acquire, removed when a release drains it. + memory_for_task: HashMap, + } + + impl SparkPool { + fn memory_free(&self) -> i64 { + self.pool_size - self.memory_for_task.values().sum::() + } + } + + /// In-process stand-in for Spark's task memory manager that models Spark 4.1.3's + /// `ExecutionMemoryPool.acquireMemory` and `releaseMemory` for this task, including the + /// per-task entry lifecycle, the 1/N and 1/(2N) share rules, and the wait loop, together + /// with `CometTaskMemoryManager` asking again when a waiter's entry is gone. Re-adding the + /// entry in the wait loop stands in for that JVM retry, so these tests do not cover the + /// retry itself. + struct StubTaskMemory { + pool: Mutex, + /// Spark's `lock`: parked acquires wait on it and every release does `notifyAll`. + lock: Condvar, + releases: AtomicUsize, + acquires: AtomicUsize, + /// The size of every acquire, in the order they reached Spark. + requests: Mutex>, + /// `CometTaskMemoryManager.used`: bytes granted and not yet released. This is what the + /// JVM checks for leaked reservations once the task's last plan has closed. + used: AtomicI64, + /// When non-zero, every n-th acquire asks Spark for only half of the requested bytes, + /// which is how a caller sees a short grant that must be rolled back. + short_every: usize, + /// When set, acquire fails outright. + fail_acquire: AtomicBool, + /// When set, release fails outright, like a JVM exception on the way back. + fail_release: AtomicBool, + /// When set, acquire panics, like a failure inside the bridge's JNI frame. + panic_acquire: AtomicBool, + /// Pauses an acquire before it takes the pool monitor, like `TaskMemoryManager` + /// spilling other consumers between two pool calls. + acquire_gate: Gate, + /// Pauses a release before it takes the pool monitor, like the JNI hop. + release_gate: Gate, + /// Pauses an acquire after Spark has granted it, before the pool sees the answer. + granted_gate: Gate, + /// Announces every trip into `lock.wait()`, so a waiter that is woken and parks again + /// announces twice. + parked: (Sender<()>, Mutex>), + } + + impl StubTaskMemory { + fn new(pool_size: usize) -> Self { + let parked = channel(); + Self { + pool: Mutex::new(SparkPool { + pool_size: pool_size as i64, + memory_for_task: HashMap::new(), + }), + lock: Condvar::new(), + releases: AtomicUsize::new(0), + acquires: AtomicUsize::new(0), + requests: Mutex::new(Vec::new()), + used: AtomicI64::new(0), + short_every: 0, + fail_acquire: AtomicBool::new(false), + fail_release: AtomicBool::new(false), + panic_acquire: AtomicBool::new(false), + acquire_gate: Gate::new(), + release_gate: Gate::new(), + granted_gate: Gate::new(), + parked: (parked.0, Mutex::new(parked.1)), + } + } + + fn short_every(mut self, n: usize) -> Self { + self.short_every = n; + self + } + + /// This task's balance in `memoryForTask`, or 0 once Spark has removed the entry. + fn outstanding(&self) -> i64 { + self.pool + .lock() + .memory_for_task + .get(&THIS_TASK) + .copied() + .unwrap_or(0) + } + + fn memory_free(&self) -> i64 { + self.pool.lock().memory_free() + } + + /// Another task of the same executor takes `bytes` from the pool, which also raises + /// `numActiveTasks` and shrinks this task's shares. + fn task_holds(&self, task: i64, bytes: i64) { + let mut pool = self.pool.lock(); + assert!( + bytes <= pool.memory_free(), + "task {task} cannot hold {bytes} bytes, only {} free", + pool.memory_free() + ); + *pool.memory_for_task.entry(task).or_insert(0) += bytes; + } + + fn other_task_holds(&self, bytes: i64) { + self.task_holds(OTHER_TASK, bytes); + } + + /// Another consumer of this task, such as the shuffle allocator, takes `bytes`. + fn sibling_consumer_holds(&self, bytes: i64) { + self.task_holds(THIS_TASK, bytes); + } + + /// A neighbour hands `bytes` back, dropping out of the active set at zero, which wakes + /// every parked acquire like Spark's notifyAll. + fn task_releases(&self, task: i64, bytes: i64) { + let mut pool = self.pool.lock(); + let balance = pool + .memory_for_task + .get_mut(&task) + .expect("task holds nothing"); + assert!(*balance >= bytes, "task {task} holds only {balance} bytes"); + *balance -= bytes; + if *balance <= 0 { + pool.memory_for_task.remove(&task); + } + self.lock.notify_all(); + } + + fn other_task_releases(&self, bytes: i64) { + self.task_releases(OTHER_TASK, bytes); + } + + /// The sibling consumer hands `bytes` of this task back, which removes the task's entry + /// if that empties it. + fn sibling_consumer_releases(&self, bytes: i64) { + self.task_releases(THIS_TASK, bytes); + } + + fn wait_parked(&self, what: &str) { + self.parked + .1 + .lock() + .recv_timeout(TEST_TIMEOUT) + .unwrap_or_else(|_| panic!("{what} never parked inside Spark")); + } + } + + impl StubTaskMemory { + /// Spark's `acquireExecutionMemory` for this task. + fn grant(&self, additional: usize) -> CometResult { + let n = self.acquires.fetch_add(1, SeqCst) + 1; + self.requests.lock().push(additional); + if self.fail_acquire.load(SeqCst) { + return Err(CometError::Internal("injected acquire failure".to_string())); + } + if self.panic_acquire.load(SeqCst) { + panic!("injected acquire panic"); + } + self.acquire_gate.pass(); + let num_bytes = if self.short_every != 0 && n.is_multiple_of(self.short_every) { + additional.div_ceil(2) + } else { + additional + } as i64; + assert!( + num_bytes > 0, + "invalid number of bytes requested: {num_bytes}" + ); + + let mut pool = self.pool.lock(); + loop { + // Spark adds the task's entry on the way in. A woken waiter whose entry a + // release has removed meanwhile gets a NoSuchElementException, and + // `CometTaskMemoryManager` then asks again, which adds it back. + if let Entry::Vacant(entry) = pool.memory_for_task.entry(THIS_TASK) { + entry.insert(0); + self.lock.notify_all(); + } + let num_active_tasks = pool.memory_for_task.len() as i64; + let cur_mem = pool.memory_for_task[&THIS_TASK]; + let max_memory_per_task = pool.pool_size / num_active_tasks; + let min_memory_per_task = pool.pool_size / (2 * num_active_tasks); + let max_to_grant = num_bytes.min((max_memory_per_task - cur_mem).max(0)); + let to_grant = max_to_grant.min(pool.memory_free()); + if to_grant < num_bytes && cur_mem + to_grant < min_memory_per_task { + let _ = self.parked.0.send(()); + let timed_out = self.lock.wait_for(&mut pool, WAIT_TIMEOUT).timed_out(); + assert!( + !timed_out, + "deadlock: acquire of {num_bytes} bytes waited {WAIT_TIMEOUT:?} for \ + memory that never came" + ); + } else { + *pool.memory_for_task.get_mut(&THIS_TASK).unwrap() += to_grant; + return Ok(to_grant); + } + } + } + + /// Spark's `releaseExecutionMemory` for this task. + fn hand_back(&self, size: usize) -> CometResult<()> { + if self.fail_release.load(SeqCst) { + return Err(CometError::Internal("injected release failure".to_string())); + } + self.release_gate.pass(); + let mut pool = self.pool.lock(); + // Spark only warns and clamps here; the pool must never hand back more than the + // task holds, so the stub makes that a hard failure. + let cur_mem = pool.memory_for_task.get(&THIS_TASK).copied().unwrap_or(0); + assert!( + cur_mem >= size as i64, + "released {size} bytes with only {cur_mem} outstanding" + ); + if let Some(balance) = pool.memory_for_task.get_mut(&THIS_TASK) { + *balance -= size as i64; + if *balance <= 0 { + pool.memory_for_task.remove(&THIS_TASK); + } + } + self.releases.fetch_add(1, SeqCst); + self.lock.notify_all(); + Ok(()) + } + } + + /// Keeps the count `CometTaskMemoryManager` keeps on top of Spark's balance. Like the JVM, + /// a release calls Spark first and moves the count only once it has answered. + impl SparkMemoryManager for Arc { + fn acquire(&self, additional: usize) -> CometResult { + let granted = self.grant(additional)?; + self.used.fetch_add(granted, SeqCst); + self.granted_gate.pass(); + Ok(granted) + } + + fn release(&self, size: usize) -> CometResult<()> { + self.hand_back(size)?; + self.used.fetch_sub(size as i64, SeqCst); + Ok(()) + } + } + + fn pool_with(stub: &Arc, pool_size: usize) -> Arc { + Arc::new(CometFairMemoryPool::with_spark( + SparkMemory::with_manager(Box::new(Arc::clone(stub)), THIS_TASK), + pool_size, + )) + } + + #[test] + fn try_grow_refusal_with_overcommit_outstanding_hands_back_the_full_grant() { + let fake = FakeSpark::with(100); + let fair = Arc::new(CometFairMemoryPool::with_spark(fake.memory(), 10_000)); + let pool: Arc = Arc::clone(&fair) as _; + let reservation = MemoryConsumer::new("consumer").register(&pool); + + // Spark grants 100 of the 150 byte grow and 50 becomes overcommit. + reservation.grow(150); + assert_eq!(fake.held(), 100); + assert_eq!(pool.reserved(), 150); + assert_eq!(fair.overcommit(), 50); + + // 20 bytes of new headroom, not enough for the 5 byte request plus the 50 owed. Spark + // hands over more than the request itself, and all of it must go back. + fake.set_limit(120); + let err = reservation.try_grow(5).unwrap_err(); + assert!(err.to_string().contains("only got"), "{err}"); + + assert_eq!(pool.reserved(), 150); + assert_eq!(fake.held(), 100); + assert_eq!(fake.released(), vec![20]); + assert_eq!( + fair.overcommit(), + 50, + "a refused try_grow repays no overcommit" + ); + } + + /// A try_grow refused while another one in flight is asking Spark for the debt reports the + /// overcommit outstanding, not only the part it asked for itself. + #[test] + fn try_grow_refusal_reports_debt_another_try_grow_is_asking_for() { + let fake = FakeSpark::with(100); + let pool: Arc = + Arc::new(CometFairMemoryPool::with_spark(fake.memory(), 10_000)); + let first = MemoryConsumer::new("first").register(&pool); + let second = MemoryConsumer::new("second").register(&pool); + // Spark grants 100 of the 150 bytes, so 50 are owed and Spark is full. + first.grow(150); + + let (err_tx, err_rx) = channel(); + fake.during_next_acquire(move || { + let _ = err_tx.send(second.try_grow(1).unwrap_err().to_string()); + }); + assert!(first.try_grow(5).is_err()); + let err = err_rx.recv().unwrap(); + assert!(err.contains("plus 0 of 50 bytes overcommitted"), "{err}"); + } #[test] fn grow_past_the_fair_limit_is_recorded_and_refuses_the_next_try_grow() { @@ -225,6 +698,1030 @@ mod tests { assert_eq!(fake.released(), vec![100]); } + #[test] + fn grow_and_shrink_update_pool_and_spark_accounting() { + let stub = Arc::new(StubTaskMemory::new(GIB)); + let pool = pool_with(&stub, 1_000); + let res = MemoryConsumer::new("consumer").register(&pool); + + res.try_grow(600).unwrap(); + assert_eq!(pool.reserved(), 600); + assert_eq!(stub.outstanding(), 600); + + res.shrink(200); + assert_eq!(pool.reserved(), 400); + assert_eq!(stub.outstanding(), 400); + + res.free(); + assert_eq!(pool.reserved(), 0); + assert_eq!(stub.outstanding(), 0); + } + + #[test] + fn try_grow_beyond_fair_limit_fails_without_calling_spark() { + let stub = Arc::new(StubTaskMemory::new(GIB)); + let pool = pool_with(&stub, 1_000); + let res = MemoryConsumer::new("consumer").register(&pool); + + res.try_grow(600).unwrap(); + let acquires_before = stub.acquires.load(SeqCst); + let err = res.try_grow(500).unwrap_err(); + assert!(err.to_string().contains("fair limit"), "{err}"); + assert_eq!(pool.reserved(), 600); + assert_eq!( + stub.acquires.load(SeqCst), + acquires_before, + "over-limit grow must be rejected before reaching Spark" + ); + res.free(); + } + + #[test] + fn fair_limit_shrinks_as_consumers_register() { + let stub = Arc::new(StubTaskMemory::new(GIB)); + let pool = pool_with(&stub, 1_000); + let first = MemoryConsumer::new("first").register(&pool); + + first.try_grow(600).unwrap(); + + // A second consumer halves the fair limit, so the pool is now over it. + let second = MemoryConsumer::new("second").register(&pool); + let err = first.try_grow(1).unwrap_err(); + assert!(err.to_string().contains("fair limit"), "{err}"); + + drop(second); + first.try_grow(1).unwrap(); + first.free(); + } + + #[test] + fn short_grant_is_released_and_reported_as_error() { + let stub = Arc::new(StubTaskMemory::new(GIB).short_every(1)); + let pool = pool_with(&stub, 1_000); + let res = MemoryConsumer::new("consumer").register(&pool); + + let err = res.try_grow(100).unwrap_err(); + assert!(err.to_string().contains("only got"), "{err}"); + assert_eq!(pool.reserved(), 0); + assert_eq!(stub.outstanding(), 0, "partial grant must be handed back"); + } + + #[test] + fn acquire_failure_leaves_accounting_unchanged() { + let stub = Arc::new(StubTaskMemory::new(GIB)); + let pool = pool_with(&stub, 1_000); + let res = MemoryConsumer::new("consumer").register(&pool); + + stub.fail_acquire.store(true, SeqCst); + assert!(res.try_grow(100).is_err()); + assert_eq!(pool.reserved(), 0); + assert_eq!(stub.outstanding(), 0); + } + + #[test] + fn zero_sized_grow_does_not_call_spark() { + let stub = Arc::new(StubTaskMemory::new(GIB)); + let pool = pool_with(&stub, 1_000); + let res = MemoryConsumer::new("consumer").register(&pool); + + let acquires_before = stub.acquires.load(SeqCst); + pool.try_grow(&res, 0).unwrap(); + pool.shrink(&res, 0); + assert_eq!(stub.acquires.load(SeqCst), acquires_before); + assert_eq!(stub.outstanding(), 0); + } + + #[test] + #[should_panic(expected = "Failed to release")] + fn shrinking_more_than_tracked_panics() { + let stub = Arc::new(StubTaskMemory::new(GIB)); + let pool = pool_with(&stub, 1_000); + let res = MemoryConsumer::new("consumer").register(&pool); + + pool.shrink(&res, 100); + } + + /// A panic escaping the bridge's acquire must propagate, but it must not leave the + /// optimistically reserved bytes behind, or the task-shared pool would be poisoned for + /// every other consumer. + #[test] + fn panicking_acquire_rolls_back_the_reservation() { + let stub = Arc::new(StubTaskMemory::new(GIB)); + let pool = pool_with(&stub, 1_000); + let res = MemoryConsumer::new("consumer").register(&pool); + res.try_grow(100).unwrap(); + + stub.panic_acquire.store(true, SeqCst); + let panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| res.try_grow(600))); + assert!(panic.is_err(), "the bridge panic must propagate"); + assert_eq!( + pool.reserved(), + 100, + "panicked grow left phantom bytes behind" + ); + + stub.panic_acquire.store(false, SeqCst); + res.try_grow(600).unwrap(); + assert_eq!(pool.reserved(), 700); + res.free(); + assert_eq!(stub.outstanding(), 0); + } + + /// `grow` cannot refuse, but a panic in the bridge's acquire still propagates, and the bytes + /// it charged before the call must come off the pool's total and the consumer's share. + #[test] + fn panicking_acquire_in_grow_rolls_back_the_reservation() { + let stub = Arc::new(StubTaskMemory::new(GIB)); + let pool = pool_with(&stub, 1_000); + let res = MemoryConsumer::new("consumer").register(&pool); + res.grow(100); + + stub.panic_acquire.store(true, SeqCst); + let panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| res.grow(600))); + assert!(panic.is_err(), "the bridge panic must propagate"); + assert_eq!( + pool.reserved(), + 100, + "panicked grow left phantom bytes behind" + ); + assert_eq!(res.size(), 100); + + // The consumer's share is back to 100 bytes, so its 1000 byte fair limit admits 900. + stub.panic_acquire.store(false, SeqCst); + res.try_grow(900) + .expect("the consumer's share still holds the panicked grow"); + assert_eq!(pool.reserved(), 1_000); + res.free(); + assert_eq!(stub.outstanding(), 0); + } + + /// Many threads hammering grow/shrink must keep the pool's accounting and the Spark-side + /// balance consistent, including through fairness rejections and partial-grant rollbacks. + #[test] + fn concurrent_grow_and_shrink_keep_accounting_consistent() { + const THREADS: usize = 8; + const ITERS: usize = 500; + const POOL_SIZE: usize = 80_000; + + let stub = Arc::new(StubTaskMemory::new(GIB).short_every(7)); + let pool = pool_with(&stub, POOL_SIZE); + + // Register every consumer up front so the fair limit stays fixed while threads run. + let reservations: Vec<_> = (0..THREADS) + .map(|t| MemoryConsumer::new(format!("consumer-{t}")).register(&pool)) + .collect(); + let barrier = Arc::new(Barrier::new(THREADS)); + + let handles: Vec<_> = reservations + .into_iter() + .enumerate() + .map(|(t, res)| { + let pool = Arc::clone(&pool); + let barrier = Arc::clone(&barrier); + thread::spawn(move || { + for i in 0..ITERS { + let size = 1 + (i * 37 + t * 101) % 509; + // Fairness rejections and short grants are expected; only the + // accounting invariants below must hold. + let _ = res.try_grow(size); + if i % 3 == 0 && res.size() > 0 { + res.shrink(res.size() / 2 + 1); + } + // Occasional full drains push the task's balance toward zero while + // other threads still have acquires in flight. + if (i + t) % 41 == 0 { + res.free(); + } + // Each consumer has one reservation, so its size is the consumer's + // share of the pool. + assert!( + res.size() <= POOL_SIZE / THREADS, + "consumer exceeded its fair limit" + ); + assert!(pool.reserved() <= POOL_SIZE, "pool exceeded its size"); + } + // Keep all consumers registered until every thread stops growing, so the + // fair-limit assertion above stays valid for the whole run. + barrier.wait(); + res.free(); + }) + }) + .collect(); + + for handle in handles { + handle.join().unwrap(); + } + + assert_eq!(pool.reserved(), 0, "pool still tracks bytes after quiesce"); + assert_eq!( + stub.outstanding(), + 0, + "Spark-side bytes leaked or double-released" + ); + assert_eq!( + stub.used.load(SeqCst), + stub.outstanding(), + "the JVM's count drifted from Spark's balance" + ); + } + + /// A thread stuck inside the blocking acquire call must not prevent another thread from + /// releasing memory: the release path cannot wait on any lock held across that call. + #[test] + fn shrink_is_not_blocked_by_a_slow_acquire_on_another_thread() { + let stub = Arc::new(StubTaskMemory::new(GIB)); + let pool = pool_with(&stub, 1_000_000); + + let holder = MemoryConsumer::new("holder").register(&pool); + let grower = MemoryConsumer::new("grower").register(&pool); + holder.try_grow(1_000).unwrap(); + + stub.acquire_gate.arm(); + let grower_thread = thread::spawn(move || { + // The result is irrelevant; the test only needs this acquire to be in flight. + let _ = grower.try_grow(500); + grower.free(); + }); + stub.acquire_gate.wait_entered("grower"); + + let (done_tx, done_rx) = channel(); + let releaser_thread = thread::spawn(move || { + holder.free(); + let _ = done_tx.send(()); + }); + let released = done_rx.recv_timeout(TEST_TIMEOUT); + + // Open the gate before asserting so no thread stays parked if the assertion fails. + stub.acquire_gate.disarm(); + stub.acquire_gate.open(); + releaser_thread.join().unwrap(); + grower_thread.join().unwrap(); + assert!( + released.is_ok(), + "release was blocked behind an in-flight acquire" + ); + } + + /// Spark can park an acquire for as long as other tasks hold the memory it needs. Reading + /// the pool's usage, as the memory usage logger does from its own thread, must return at + /// once meanwhile, so no lock that `reserved` takes may be held across the Spark call. + #[test] + fn reserved_is_not_blocked_by_an_acquire_parked_in_spark() { + let stub = Arc::new(StubTaskMemory::new(100)); + let pool = pool_with(&stub, 1_000_000); + + let holder = MemoryConsumer::new("holder").register(&pool); + let grower = MemoryConsumer::new("grower").register(&pool); + holder.try_grow(10).unwrap(); + // Fill the pool so Spark parks the grower below its minimum share until the holder + // frees, which only happens after the read below. + stub.other_task_holds(stub.memory_free()); + + let grower_thread = thread::spawn(move || { + grower.try_grow(10).unwrap(); + grower.free(); + }); + stub.wait_parked("grower"); + + let (reserved_tx, reserved_rx) = channel(); + let reader_pool = Arc::clone(&pool); + let reader_thread = thread::spawn(move || { + let _ = reserved_tx.send(reader_pool.reserved()); + }); + // Keep this wait well below the stub's WAIT_TIMEOUT so a blocked read shows up as a + // timeout. Once the stub gives up, the grower's charge rolls back and a blocked read + // returns 10 instead, which fails as a less telling count mismatch. + let reserved_during = reserved_rx.recv_timeout(Duration::from_secs(1)); + + // Freeing the holder wakes the grower. If the lock were held across the Spark call, the + // free would itself wait until the stub gives up, which still bounds the run. + holder.free(); + let grower_result = grower_thread.join(); + reader_thread.join().unwrap(); + assert_eq!( + reserved_during, + Ok(20), + "reserved() must return at once and count the holder's 10 bytes plus the grower's \ + 10 bytes in flight" + ); + grower_result.expect("parked acquire failed after the holder freed"); + assert_eq!(pool.reserved(), 0); + } + + /// Two threads of one task: the pool is full and this task sits below its 1/(2N) minimum + /// share, so Spark parks the second acquire until memory is freed. The first thread then + /// frees everything it holds. The release must go through in full and be what wakes the + /// waiter. It takes the task's balance to zero, so Spark drops the task's entry and the + /// waiter fails on waking; the bridge asks again and the freed bytes are granted. + #[test] + fn full_release_wakes_an_acquire_parked_below_its_minimum_share() { + let stub = Arc::new(StubTaskMemory::new(100)); + let pool = pool_with(&stub, 1_000_000); + + let holder = MemoryConsumer::new("holder").register(&pool); + let grower = MemoryConsumer::new("grower").register(&pool); + holder.try_grow(10).unwrap(); + // Fill the pool: two active tasks, this one at 10 of a 25 byte minimum share. + stub.other_task_holds(stub.memory_free()); + + let (grower_done_tx, grower_done_rx) = channel(); + let grower_thread = thread::spawn(move || { + grower.try_grow(10).unwrap(); + grower.free(); + let _ = grower_done_tx.send(()); + }); + stub.wait_parked("grower"); + + let (holder_done_tx, holder_done_rx) = channel(); + let holder_thread = thread::spawn(move || { + holder.free(); + let _ = holder_done_tx.send(()); + }); + + assert!( + holder_done_rx.recv_timeout(TEST_TIMEOUT).is_ok(), + "full release deadlocked behind the parked acquire" + ); + assert!( + grower_done_rx.recv_timeout(TEST_TIMEOUT).is_ok(), + "parked acquire never completed after the release" + ); + holder_thread.join().unwrap(); + grower_thread + .join() + .expect("parked acquire crashed after the full release"); + assert_eq!(pool.reserved(), 0); + assert_eq!(stub.outstanding(), 0); + } + + /// Modeled on Spark 4.1.3's `ExecutionMemoryPool`: a 1 GiB execution pool, another task + /// holding 924 MiB, this task holding 100 MiB, and a second native consumer of this task + /// asking for 100 MiB. Spark parks it below the 256 MiB minimum share. When the holder frees + /// its 100 MiB, Spark computes toGrant = min(request, memoryFree); anything less than the + /// full 100 MiB leaves toGrant < request with curMem + toGrant still under 256 MiB, and the + /// waiter sleeps again with nobody left to wake it. + #[test] + fn min_share_wait_is_granted_after_the_holder_frees_its_memory_in_full() { + let stub = Arc::new(StubTaskMemory::new(GIB)); + let pool = pool_with(&stub, GIB); + + let holder = MemoryConsumer::new("holder").register(&pool); + let grower = MemoryConsumer::new("grower").register(&pool); + holder.try_grow(100 * MIB).unwrap(); + // The neighbour takes the rest of the pool, 924 MiB. + stub.other_task_holds(stub.memory_free()); + assert_eq!(stub.memory_free(), 0); + + let (grower_done_tx, grower_done_rx) = channel(); + let grower_thread = thread::spawn(move || { + grower.try_grow(100 * MIB).unwrap(); + grower.free(); + let _ = grower_done_tx.send(()); + }); + stub.wait_parked("grower"); + + let holder_thread = thread::spawn(move || holder.free()); + holder_thread.join().unwrap(); + + assert!( + grower_done_rx.recv_timeout(TEST_TIMEOUT).is_ok(), + "acquire stayed parked after the holder freed 100 MiB" + ); + grower_thread + .join() + .expect("parked acquire crashed or timed out"); + assert_eq!(pool.reserved(), 0); + assert_eq!(stub.outstanding(), 0); + } + + /// A release that is already on its way to Spark when a new acquire arrives and parks: the + /// release removes the task's entry from under the waiter, and the bridge asks again. + /// Reproduced by holding the release at the JNI hop, starting a grow that parks below its + /// minimum share, then letting the release land. + #[test] + fn late_acquire_survives_a_release_already_on_its_way() { + let stub = Arc::new(StubTaskMemory::new(100)); + let pool = pool_with(&stub, 1_000_000); + + let holder = MemoryConsumer::new("holder").register(&pool); + let grower = MemoryConsumer::new("grower").register(&pool); + holder.try_grow(10).unwrap(); + stub.other_task_holds(stub.memory_free()); + + // The holder's full release is planned and dispatched, then held before Spark sees it. + stub.release_gate.arm(); + let holder_thread = thread::spawn(move || holder.free()); + stub.release_gate.wait_entered("holder release"); + + // Only now does the grower arrive; the pool is full so it parks inside Spark. + let grower_thread = thread::spawn(move || { + grower.try_grow(10).unwrap(); + grower.free(); + }); + stub.wait_parked("grower"); + + stub.release_gate.disarm(); + stub.release_gate.open(); + holder_thread.join().unwrap(); + grower_thread + .join() + .expect("late acquire crashed when the earlier release landed"); + assert_eq!(pool.reserved(), 0); + assert_eq!(stub.outstanding(), 0); + } + + /// A short grant is handed back whole, and that release can land while a sibling of this + /// task is parked inside Spark. Here the first consumer's rollback is held at the JNI hop, a + /// third task then leaves the pool (which raises this task's minimum share), and a second + /// consumer's acquire arrives and parks. The rollback takes the task's balance to zero and + /// wakes it, and the acquire gets the bytes once the bridge has asked again. + #[test] + fn short_grant_rollback_landing_under_a_parked_sibling_wakes_it() { + let stub = Arc::new(StubTaskMemory::new(100)); + let pool = pool_with(&stub, 1_000_000); + + let first = MemoryConsumer::new("first").register(&pool); + let second = MemoryConsumer::new("second").register(&pool); + // Three active tasks: 17 free with a 16 byte minimum share, so a 30 byte request gets + // a short grant of 17 rather than parking. + stub.task_holds(OTHER_TASK, 82); + stub.task_holds(THIRD_TASK, 1); + + stub.release_gate.arm(); + let first_thread = thread::spawn(move || { + let result = first.try_grow(30); + assert!(result.is_err(), "30 bytes cannot fit into 17 free bytes"); + first.free(); + }); + stub.release_gate + .wait_entered("first consumer's rollback release"); + + // The third task leaves: two active tasks now, a 25 byte minimum share, 1 byte free. + stub.task_releases(THIRD_TASK, 1); + let second_thread = thread::spawn(move || { + second.try_grow(10).unwrap(); + second.free(); + }); + stub.wait_parked("second acquire"); + + stub.release_gate.disarm(); + stub.release_gate.open(); + first_thread.join().unwrap(); + second_thread + .join() + .expect("second acquire crashed when the rollback release landed"); + assert_eq!(pool.reserved(), 0); + assert_eq!(stub.outstanding(), 0); + } + + /// A pool that never grows never touches Spark. + #[test] + fn pool_that_never_grows_never_calls_spark() { + let stub = Arc::new(StubTaskMemory::new(GIB)); + let pool = pool_with(&stub, 1_000); + let res = MemoryConsumer::new("consumer").register(&pool); + drop(res); + drop(pool); + assert_eq!(stub.acquires.load(SeqCst), 0); + assert_eq!(stub.releases.load(SeqCst), 0); + } + + /// Spark's shuffle allocator is another off-heap consumer of the same task. When it frees + /// its last page while an acquire of this pool is parked, Spark drops the task's entry, and + /// the acquire is granted once the bridge has asked again. + #[test] + fn sibling_consumer_freeing_its_last_page_under_a_parked_acquire() { + let stub = Arc::new(StubTaskMemory::new(100)); + let pool = pool_with(&stub, 1_000_000); + let grower = MemoryConsumer::new("grower").register(&pool); + + // A sibling consumer holds 10 bytes of this task and a neighbour fills the rest, so a + // 10 byte request parks below the 25 byte minimum share. + stub.sibling_consumer_holds(10); + stub.other_task_holds(stub.memory_free()); + let grower_thread = thread::spawn(move || { + grower.try_grow(10).unwrap(); + grower.free(); + }); + stub.wait_parked("grower"); + + stub.sibling_consumer_releases(10); + grower_thread + .join() + .expect("parked acquire crashed when a sibling consumer freed its last page"); + assert_eq!(pool.reserved(), 0); + assert_eq!(stub.outstanding(), 0); + } + + /// One consumer's acquire is parked inside Spark. A second consumer's acquire that Spark can + /// grant at once must not queue behind it on the pool's lock. This covers the Rust pool's lock + /// only: in the JVM, `TaskMemoryManager` serializes a task's acquires on its own monitor, + /// which the parked acquire keeps while it waits. + #[test] + fn second_acquire_does_not_queue_behind_a_parked_one() { + let stub = Arc::new(StubTaskMemory::new(100)); + let pool = pool_with(&stub, 1_000_000); + + let first = MemoryConsumer::new("first").register(&pool); + let second = MemoryConsumer::new("second").register(&pool); + // 4 bytes free with a 25 byte minimum share: a 20 byte request parks, a 4 byte one + // is granted outright. + stub.other_task_holds(stub.memory_free() - 4); + + let first_thread = thread::spawn(move || { + let result = first.try_grow(20); + first.free(); + result + }); + stub.wait_parked("first acquire"); + + let (second_done_tx, second_done_rx) = channel(); + let second_thread = thread::spawn(move || { + let result = second.try_grow(4); + second.free(); + let _ = second_done_tx.send(result); + }); + let second_result = second_done_rx.recv_timeout(Duration::from_secs(2)); + + // Room for the parked request, released before asserting so no thread stays parked. + stub.other_task_releases(50); + let first_result = first_thread.join().expect("parked first acquire crashed"); + first_result.expect("parked first acquire was not granted once memory was freed"); + second_thread.join().unwrap(); + second_result + .expect("second acquire queued behind the parked first acquire") + .expect("second acquire was not granted"); + assert_eq!(pool.reserved(), 0); + assert_eq!(stub.outstanding(), 0); + } + + /// A pool whose `overcommit` a test can read, connected to `stub`. + fn fair_pool_with(stub: &Arc, pool_size: usize) -> Arc { + Arc::new(CometFairMemoryPool::with_spark( + SparkMemory::with_manager(Box::new(Arc::clone(stub)), THIS_TASK), + pool_size, + )) + } + + /// Asks for `size` bytes through `try_grow` or, with `infallible`, through `grow`, and checks + /// that Spark granted all of them in a single call without parking. + fn assert_granted_in_one_call( + stub: &StubTaskMemory, + fair: &CometFairMemoryPool, + res: &MemoryReservation, + size: usize, + infallible: bool, + ) { + let acquires_before = stub.acquires.load(SeqCst); + let outstanding_before = stub.outstanding(); + if infallible { + res.grow(size); + } else { + res.try_grow(size) + .unwrap_or_else(|e| panic!("a {size} byte request Spark can grant failed: {e}")); + } + assert_eq!( + stub.acquires.load(SeqCst), + acquires_before + 1, + "one Spark call" + ); + assert_eq!( + stub.outstanding(), + outstanding_before + size as i64, + "Spark granted the request in full" + ); + assert_eq!(fair.overcommit(), 0, "nothing carried as overcommit"); + assert!(stub.parked.1.lock().try_recv().is_err(), "nothing parked"); + } + + /// Another task holds 90 of Spark's 100 bytes and this task holds nothing, so a 10 byte + /// request is exactly what is free and is under the task's 25 byte minimum share. Spark + /// grants it at once, and the pool must ask for nothing besides it. + fn exact_fit_below_the_minimum_share(infallible: bool) { + let stub = Arc::new(StubTaskMemory::new(100)); + stub.other_task_holds(90); + let fair = fair_pool_with(&stub, 1_000); + let pool: Arc = Arc::clone(&fair) as _; + let res = MemoryConsumer::new("consumer").register(&pool); + + assert_granted_in_one_call(&stub, &fair, &res, 10, infallible); + assert_eq!(pool.reserved(), 10); + res.free(); + assert_eq!(stub.outstanding(), 0); + } + + #[test] + fn exact_fit_try_grow_below_the_minimum_share_is_granted_in_one_call() { + exact_fit_below_the_minimum_share(false); + } + + #[test] + fn exact_fit_grow_below_the_minimum_share_is_granted_in_one_call() { + exact_fit_below_the_minimum_share(true); + } + + /// The same exact fit for a request that reaches the task's minimum share: another task + /// holds 70 bytes, and 30 are granted in full rather than short. + fn exact_fit_at_the_minimum_share(infallible: bool) { + let stub = Arc::new(StubTaskMemory::new(100)); + stub.other_task_holds(70); + let fair = fair_pool_with(&stub, 1_000); + let pool: Arc = Arc::clone(&fair) as _; + let res = MemoryConsumer::new("consumer").register(&pool); + + assert_granted_in_one_call(&stub, &fair, &res, 30, infallible); + assert_eq!(pool.reserved(), 30); + res.free(); + assert_eq!(stub.outstanding(), 0); + } + + #[test] + fn exact_fit_try_grow_at_the_minimum_share_is_granted_in_full() { + exact_fit_at_the_minimum_share(false); + } + + #[test] + fn exact_fit_grow_at_the_minimum_share_is_granted_in_full() { + exact_fit_at_the_minimum_share(true); + } + + /// Once the pool has handed everything back, it holds nothing from Spark, so the next exact + /// fit sees the same free bytes as the first one and is granted in one call too. + fn exact_fit_after_a_full_release(infallible: bool) { + let stub = Arc::new(StubTaskMemory::new(100)); + let fair = fair_pool_with(&stub, 1_000); + let pool: Arc = Arc::clone(&fair) as _; + let res = MemoryConsumer::new("consumer").register(&pool); + res.try_grow(10).unwrap(); + res.free(); + assert_eq!( + stub.outstanding(), + 0, + "a full release leaves nothing with Spark" + ); + + stub.other_task_holds(90); + assert_granted_in_one_call(&stub, &fair, &res, 10, infallible); + res.free(); + assert_eq!(stub.outstanding(), 0); + } + + #[test] + fn exact_fit_try_grow_after_a_full_release_is_granted_in_one_call() { + exact_fit_after_a_full_release(false); + } + + #[test] + fn exact_fit_grow_after_a_full_release_is_granted_in_one_call() { + exact_fit_after_a_full_release(true); + } + + /// A shrink hands its bytes back to Spark before the pool stops counting them. While that + /// release is still on its way, a grow of the same consumer that only fits if the shrunk + /// bytes were free is refused at the fair limit without a JVM call, instead of reaching + /// Spark ahead of the release and coming back with a short grant. + #[test] + fn grow_racing_a_shrink_is_not_admitted_on_bytes_the_jvm_still_holds() { + let stub = Arc::new(StubTaskMemory::new(100)); + // Two consumers: a 100 byte fair limit against a task whose Spark share is 100 bytes. + let pool = pool_with(&stub, 200); + let holder = MemoryConsumer::new("holder").register(&pool); + // A sibling reservation draws on the holder's share. + let grower = holder.new_empty(); + let other = MemoryConsumer::new("other").register(&pool); + holder.try_grow(100).unwrap(); + assert_eq!(stub.outstanding(), 100, "the task sits at its Spark share"); + + stub.release_gate.arm(); + let holder_thread = thread::spawn(move || { + holder.shrink(50); + holder + }); + stub.release_gate.wait_entered("holder release"); + + let acquires_before = stub.acquires.load(SeqCst); + let result = grower.try_grow(50); + let acquires_during = stub.acquires.load(SeqCst); + let reserved_during = pool.reserved(); + + // Let the release land before asserting so no thread stays parked on failure. + stub.release_gate.disarm(); + stub.release_gate.open(); + let holder = holder_thread.join().unwrap(); + + let err = result.expect_err("the shrunk bytes are not free until Spark has them"); + assert!(err.to_string().contains("fair limit"), "{err}"); + assert_eq!( + acquires_during, acquires_before, + "grow reached Spark on bytes the JVM still held" + ); + assert_eq!( + reserved_during, 100, + "bytes stay charged until the release lands" + ); + + grower + .try_grow(50) + .expect("the same grow fits once the release has landed"); + assert_eq!(pool.reserved(), 100); + assert_eq!(stub.outstanding(), 100, "50 held and the 50 granted"); + holder.free(); + grower.free(); + drop(holder); + drop(grower); + drop(other); + drop(pool); + assert_eq!(stub.outstanding(), 0); + } + + /// A short grant is rolled back by handing the granted bytes to Spark. Until that + /// release lands, the pool keeps charging them, so a grow of the same consumer on another + /// thread cannot be admitted on bytes Spark still holds for this task. + #[test] + fn short_grant_rollback_keeps_the_bytes_charged_until_the_jvm_has_them_back() { + let stub = Arc::new(StubTaskMemory::new(GIB).short_every(2)); + // Two consumers: a 100 byte fair limit. + let pool = pool_with(&stub, 200); + let first = MemoryConsumer::new("first").register(&pool); + // A sibling reservation draws on the first consumer's share. + let second = first.new_empty(); + let other = MemoryConsumer::new("other").register(&pool); + // The stub halves every second acquire. This 1 byte grow takes the first, so the 100 + // byte request below is granted 50 and the 60 byte grow after the rollback, the third, + // in full. + first.try_grow(1).unwrap(); + first.free(); + + stub.release_gate.arm(); + let first_thread = thread::spawn(move || { + let result = first.try_grow(100); + first.free(); + result + }); + stub.release_gate.wait_entered("short grant rollback"); + + let reserved_during = pool.reserved(); + let acquires_before = stub.acquires.load(SeqCst); + let result = second.try_grow(60); + let acquires_during = stub.acquires.load(SeqCst); + + stub.release_gate.disarm(); + stub.release_gate.open(); + let first_result = first_thread.join().unwrap(); + + let err = first_result.expect_err("100 bytes were requested and 50 granted"); + assert!(err.to_string().contains("only got"), "{err}"); + assert_eq!( + reserved_during, 50, + "the granted bytes stay charged during the rollback" + ); + let err = result.expect_err("60 bytes do not fit beside 50 still charged"); + assert!(err.to_string().contains("fair limit"), "{err}"); + assert_eq!( + acquires_during, acquires_before, + "grow reached Spark on bytes the JVM still held" + ); + + assert_eq!( + pool.reserved(), + 0, + "the rollback settles once Spark has the bytes" + ); + assert_eq!(stub.outstanding(), 0); + second + .try_grow(60) + .expect("the same grow fits once the rollback has landed"); + assert_eq!(pool.reserved(), 60); + second.free(); + drop(second); + drop(other); + drop(pool); + assert_eq!(stub.outstanding(), 0); + } + + /// A grow is charged to the pool when it passes the limits, before Spark answers it. While + /// it is in flight, another consumer's grow that would push the total past the pool size is + /// refused, which is what keeps two grows from jointly exceeding the pool. The charge is + /// rolled back if Spark does not back it. + #[test] + fn a_grow_in_flight_counts_against_the_pool_limit_until_it_settles() { + let stub = Arc::new(StubTaskMemory::new(GIB)); + let pool = pool_with(&stub, 100); + let grower = MemoryConsumer::new("grower").register(&pool); + + // Alone, the grower's share is the whole pool. + stub.acquire_gate.arm(); + let grower_thread = thread::spawn(move || { + grower.try_grow(60).unwrap(); + grower + }); + stub.acquire_gate.wait_entered("grower acquire"); + + // A second consumer halves the shares. The 60 bytes in flight still count for the pool. + let other = MemoryConsumer::new("other").register(&pool); + let reserved_during = pool.reserved(); + let acquires_before = stub.acquires.load(SeqCst); + let result = other.try_grow(50); + let acquires_during = stub.acquires.load(SeqCst); + + stub.acquire_gate.disarm(); + stub.acquire_gate.open(); + let grower = grower_thread.join().unwrap(); + + assert_eq!(reserved_during, 60, "the grow in flight is already charged"); + let err = result.expect_err("60 in flight plus 50 exceeds the 100 byte pool"); + assert!(err.to_string().contains("pool limit"), "{err}"); + assert_eq!( + acquires_during, acquires_before, + "the refused grow must not reach Spark" + ); + + assert_eq!(pool.reserved(), 60); + other + .try_grow(40) + .expect("40 bytes fit beside the settled 60"); + assert_eq!(pool.reserved(), 100); + drop(grower); + drop(other); + drop(pool); + assert_eq!(stub.outstanding(), 0); + } + + /// The rollback of a short grant hands the granted bytes back, and that release can fail. + /// Spark then holds the bytes until the task ends, so the pool keeps them charged, and the + /// caller still gets the short grant error so a spillable operator spills. + #[test] + fn failed_short_grant_rollback_keeps_the_bytes_charged_and_reports_a_short_grant() { + // Every acquire is granted half. + let stub = Arc::new(StubTaskMemory::new(GIB).short_every(1)); + // Two consumers: a 100 byte fair limit. + let pool = pool_with(&stub, 200); + let first = MemoryConsumer::new("first").register(&pool); + // A sibling reservation draws on the first consumer's share. + let second = first.new_empty(); + let other = MemoryConsumer::new("other").register(&pool); + stub.fail_release.store(true, SeqCst); + + let err = first.try_grow(100).unwrap_err(); + assert!(err.to_string().contains("only got"), "{err}"); + assert_eq!( + pool.reserved(), + 50, + "the bytes Spark still holds stay charged" + ); + assert_eq!(stub.outstanding(), 50, "the stranded grant"); + assert_eq!(stub.used.load(SeqCst), 50, "the JVM counts it as held too"); + + let acquires_before = stub.acquires.load(SeqCst); + let err = second.try_grow(51).unwrap_err(); + assert!(err.to_string().contains("fair limit"), "{err}"); + assert_eq!(stub.acquires.load(SeqCst), acquires_before); + + stub.fail_release.store(false, SeqCst); + drop(first); + drop(second); + drop(other); + drop(pool); + assert_eq!( + stub.outstanding(), + 50, + "the stranded grant stays with Spark" + ); + } + + /// A request that overflows the running total is refused like any other over-limit grow, + /// without reaching Spark. + #[test] + fn grow_that_overflows_the_total_is_refused_without_calling_spark() { + let stub = Arc::new(StubTaskMemory::new(GIB)); + let pool = pool_with(&stub, 1_000); + let res = MemoryConsumer::new("consumer").register(&pool); + res.try_grow(1).unwrap(); + + let acquires_before = stub.acquires.load(SeqCst); + let err = res.try_grow(usize::MAX).unwrap_err(); + assert!(err.to_string().contains("fair limit"), "{err}"); + assert_eq!(stub.acquires.load(SeqCst), acquires_before); + assert_eq!(pool.reserved(), 1); + res.free(); + } + + /// The window that settling after the JVM call leaves open. A grow that fits under the fair + /// limit without the shrunk bytes still reaches Spark ahead of the shrink's release, and at + /// the task's Spark share it comes back with a zero grant. This pins the behaviour rather + /// than changing it. + #[test] + fn grow_racing_a_shrink_gets_a_short_grant_at_the_spark_share() { + let stub = Arc::new(StubTaskMemory::new(100)); + // Two consumers: a 100 byte fair limit against a task whose Spark share is 100 bytes. + let pool = pool_with(&stub, 200); + let holder = MemoryConsumer::new("holder").register(&pool); + let grower = MemoryConsumer::new("grower").register(&pool); + holder.try_grow(100).unwrap(); + + stub.release_gate.arm(); + let holder_thread = thread::spawn(move || { + holder.shrink(50); + holder + }); + stub.release_gate.wait_entered("holder release"); + + let acquires_before = stub.acquires.load(SeqCst); + let result = grower.try_grow(1); + let acquires_during = stub.acquires.load(SeqCst); + let reserved_during = pool.reserved(); + + stub.release_gate.disarm(); + stub.release_gate.open(); + let holder = holder_thread.join().unwrap(); + + let err = result.expect_err("the task is at its Spark share until the release lands"); + assert!(err.to_string().contains("only got 0"), "{err}"); + assert_eq!( + acquires_during, + acquires_before + 1, + "the grow reached Spark" + ); + assert_eq!(reserved_during, 100); + + grower + .try_grow(1) + .expect("the grow fits once the release has landed"); + assert_eq!(stub.outstanding(), 51, "50 held and the 1 granted"); + drop(holder); + drop(grower); + drop(pool); + assert_eq!(stub.outstanding(), 0); + } + + /// Two try_grows of one task race on the same overcommit. The first asks Spark for its byte + /// plus the one byte owed and is held after Spark grants both, before it repays the debt. + /// The second must ask for its own byte only and finish while the first is held: one byte is + /// free, and a two byte request parks below the task's minimum share with nothing left to + /// wake it once the debt is repaid. + #[test] + fn concurrent_try_grows_ask_spark_for_the_overcommit_once() { + let stub = Arc::new(StubTaskMemory::new(100)); + let fair = fair_pool_with(&stub, 100); + let pool: Arc = Arc::clone(&fair) as _; + let first = MemoryConsumer::new("first").register(&pool); + let second = MemoryConsumer::new("second").register(&pool); + // Three active tasks: 16 bytes free and a 16 byte minimum share, so a 17 byte grow is + // granted 16 and one byte is owed. + stub.task_holds(OTHER_TASK, 83); + stub.task_holds(THIRD_TASK, 1); + first.grow(17); + assert_eq!(stub.outstanding(), 16); + assert_eq!(fair.overcommit(), 1); + // Two active tasks, a 25 byte minimum share and 3 bytes free. + stub.task_releases(THIRD_TASK, 1); + stub.other_task_releases(2); + + stub.granted_gate.arm(); + let first_thread = thread::spawn(move || first.try_grow(1).map(|()| first)); + stub.granted_gate + .wait_entered("first try_grow after its grant"); + stub.granted_gate.disarm(); + + let (second_done_tx, second_done_rx) = channel(); + let second_thread = thread::spawn(move || { + let result = second.try_grow(1); + let _ = second_done_tx.send(()); + result.map(|()| second) + }); + // The first is still held at the gate. Wait well below the stub's WAIT_TIMEOUT, so a + // request parked in Spark shows up here rather than as a stub panic. + let second_done_while_first_held = + second_done_rx.recv_timeout(Duration::from_secs(2)).is_ok(); + stub.granted_gate.open(); + let requests = stub.requests.lock().clone(); + if !second_done_while_first_held { + // The first repaying the debt wakes nothing parked in Spark, so free memory to let + // both threads finish before the assertions below fail. + stub.other_task_releases(10); + } + let first = first_thread.join().unwrap().expect("first try_grow"); + let second_result = second_thread.join(); + + assert!( + second_done_while_first_held, + "the second try_grow did not finish while the first was held after its grant" + ); + assert_eq!(requests, vec![17, 2, 1]); + let second = second_result.unwrap().expect("second try_grow"); + assert_eq!(fair.overcommit(), 0); + assert_eq!(fair.spark.claimed(), 0); + assert_eq!( + stub.outstanding() as usize + fair.overcommit(), + pool.reserved(), + "Spark's grant plus the overcommit is what the pool recorded" + ); + assert_eq!(pool.reserved(), 19); + drop(first); + drop(second); + assert_eq!(stub.outstanding(), 0); + } + #[test] fn each_consumer_is_limited_to_its_own_share() { // Spark grants everything, so only the pool's own checks refuse. diff --git a/native/core/src/execution/memory_pools/spark_memory.rs b/native/core/src/execution/memory_pools/spark_memory.rs index 39a8030a572..6506141e369 100644 --- a/native/core/src/execution/memory_pools/spark_memory.rs +++ b/native/core/src/execution/memory_pools/spark_memory.rs @@ -15,13 +15,11 @@ // specific language governing permissions and limitations // under the License. -use std::sync::{ - atomic::{AtomicUsize, Ordering::Relaxed}, - Arc, -}; +use std::sync::Arc; use jni::objects::{Global, JObject}; use log::warn; +use parking_lot::Mutex; use crate::{errors::CometResult, jvm_bridge::JVMClasses}; @@ -58,34 +56,87 @@ impl SparkMemoryManager for JniMemoryManager { /// such as a spilled batch read back from disk. When Spark grants less than [`Self::acquire`] /// asked for, the shortfall is carried as overcommit rather than failing. [`Self::release`] repays /// it before returning anything to Spark, so Spark is never handed back more than it granted. -/// While any is outstanding, [`Self::try_acquire`] also asks Spark for it, so a pool refuses -/// `try_grow` until Spark can cover both the request and the debt, and the operator spills. +/// While any is outstanding, [`Self::try_acquire`] also asks Spark for the part no other call is +/// asking for, so a `try_grow` that carries debt is refused unless Spark covers both, and the +/// operator spills. +/// +/// A `try_acquire` claims the debt that no other `try_acquire` in flight is asking Spark for, +/// and succeeds once Spark covers the request plus that claim. Debt already claimed is not asked +/// for twice, and a refused claim is given up for the next `try_acquire` to take. +/// +/// A release while a claim's request is parked in Spark leaves it too large: it waits until memory +/// frees or the number of active tasks changes, and may come back short and be refused. Handing +/// the claim to Spark on release would break the invariant below whenever the claim is refused. /// /// # Invariant /// /// Every byte a pool records through this type is backed either by Spark's grant or by /// overcommit, so Spark's grant plus `overcommit` equals the bytes recorded and not yet released. /// Spark is handed back more than it granted only if a release takes less from `overcommit` than -/// it could. `CometUnifiedMemoryPool` calls in here from several threads without a lock, and -/// updates its own `used` separately, so this rests on three things: +/// it could. `CometUnifiedMemoryPool` and `CometFairMemoryPool` both call in here from several +/// threads without holding a lock, and each updates its own `used` separately, so this rests on +/// three things: /// /// - `overcommit` only grows by bytes that are being recorded in the same call. -/// - Each repayment takes its share of `overcommit` in a single atomic update, so two concurrent +/// - Each repayment takes its share of `overcommit` under the ledger's lock, so two concurrent /// calls can never repay the same debt. /// - A caller never releases more than it recorded. DataFusion guarantees this, because a shrink /// cannot exceed the reservation it comes from. pub(super) struct SparkMemory { manager: Box, - overcommit: AtomicUsize, + /// Held only for arithmetic, never across a call to Spark. + ledger: Mutex, task_attempt_id: i64, } +#[derive(Default)] +struct Ledger { + /// Bytes recorded without Spark's grant. + overcommit: usize, + /// Debt that `try_acquire` calls in flight have put into their Spark requests. A release + /// can repay that debt meanwhile, so this can exceed `overcommit`. + claimed: usize, +} + +/// Debt one `try_acquire` has put into its Spark request. Dropped without [`Claim::settle`], +/// as on a refusal, an error or a panic, it gives the debt back for another `try_acquire`. +struct Claim<'a> { + ledger: &'a Mutex, + bytes: usize, + settled: bool, +} + +impl Claim<'_> { + /// Repays what is still owed of the claim once Spark has granted it, and returns the part a + /// release repaid meanwhile, which Spark granted for nothing. + fn settle(mut self) -> usize { + self.settled = true; + if self.bytes == 0 { + return 0; + } + let mut ledger = self.ledger.lock(); + let owed = self.bytes.min(ledger.overcommit); + ledger.overcommit -= owed; + ledger.claimed -= self.bytes; + self.bytes - owed + } +} + +impl Drop for Claim<'_> { + fn drop(&mut self) { + if !self.settled && self.bytes > 0 { + self.ledger.lock().claimed -= self.bytes; + } + } +} + /// Why [`SparkMemory::try_acquire`] refused a request. #[derive(Debug, PartialEq, Eq)] pub(super) struct Refusal { - /// Outstanding overcommit that was asked for on top of the request. + /// Overcommit this call claimed and asked for on top of the request. pub(super) overcommit: usize, - /// What Spark offered before it was handed back. + /// What Spark granted toward the request and the overcommit together, which can exceed + /// the request. pub(super) granted: usize, } @@ -97,7 +148,7 @@ impl SparkMemory { pub(super) fn with_manager(manager: Box, task_attempt_id: i64) -> Self { Self { manager, - overcommit: AtomicUsize::new(0), + ledger: Mutex::default(), task_attempt_id, } } @@ -106,28 +157,46 @@ impl SparkMemory { self.task_attempt_id } - /// Acquires `size` bytes plus any outstanding overcommit, or nothing. A full grant repays the - /// overcommit; a partial one is handed back and reported as a [`Refusal`]. + /// The Spark calls themselves, without the overcommit ledger, for a short grant a pool + /// hands back itself. + pub(super) fn manager(&self) -> &dyn SparkMemoryManager { + self.manager.as_ref() + } + + /// Acquires `size` bytes plus the overcommit it claims, or nothing. A full grant repays that + /// debt; a partial one is handed back and reported as a [`Refusal`]. pub(super) fn try_acquire(&self, size: usize) -> CometResult> { - let debt = self.overcommit.load(Relaxed); - let request = size.saturating_add(debt); + let refusal = match self.try_acquire_leaving_a_short_grant(size)? { + Ok(()) => return Ok(Ok(())), + Err(refusal) => refusal, + }; + if refusal.granted > 0 { + self.manager.release(refusal.granted)?; + } + Ok(Err(refusal)) + } + + /// Like [`Self::try_acquire`], except that a short grant stays with Spark, recorded nowhere, + /// so the caller can keep charging those bytes until it hands them back through + /// [`Self::manager`]. + pub(super) fn try_acquire_leaving_a_short_grant( + &self, + size: usize, + ) -> CometResult> { + let claim = self.claim_debt(); + let request = size.saturating_add(claim.bytes); let granted = granted(request, self.ask_spark(request)?); if granted < request { - if granted > 0 { - self.manager.release(granted)?; - } return Ok(Err(Refusal { - overcommit: debt, + overcommit: claim.bytes, granted, })); } - if debt > 0 { - // A concurrent release may have repaid part of the debt since it was read, in which - // case Spark granted more than is still owed and the excess goes back. - let owed = self.repay(debt); - if owed < debt { - self.manager.release(debt - owed)?; - } + // A concurrent release may have repaid part of the claim while Spark was asked for it, + // and Spark gets that part back. + let excess = claim.settle(); + if excess > 0 { + self.manager.release(excess)?; } Ok(Ok(())) } @@ -146,7 +215,7 @@ impl SparkMemory { } }; if granted < size { - self.overcommit.fetch_add(size - granted, Relaxed); + self.ledger.lock().overcommit += size - granted; } } @@ -160,7 +229,24 @@ impl SparkMemory { } pub(super) fn overcommit(&self) -> usize { - self.overcommit.load(Relaxed) + self.ledger.lock().overcommit + } + + #[cfg(test)] + pub(super) fn claimed(&self) -> usize { + self.ledger.lock().claimed + } + + /// Claims the debt that no other `try_acquire` in flight is asking Spark for. + fn claim_debt(&self) -> Claim<'_> { + let mut ledger = self.ledger.lock(); + let bytes = ledger.overcommit.saturating_sub(ledger.claimed); + ledger.claimed += bytes; + Claim { + ledger: &self.ledger, + bytes, + settled: false, + } } /// Asks Spark for `size` bytes and returns how many it granted. @@ -173,12 +259,12 @@ impl SparkMemory { tokio::task::block_in_place(|| self.manager.acquire(size)) } - /// Takes up to `size` bytes off the overcommit in one atomic step and returns how many. + /// Takes up to `size` bytes off the overcommit and returns how many. fn repay(&self, size: usize) -> usize { - let debt = self - .overcommit - .update(Relaxed, Relaxed, |debt| debt.saturating_sub(size)); - debt.min(size) + let mut ledger = self.ledger.lock(); + let owed = size.min(ledger.overcommit); + ledger.overcommit -= owed; + owed } } @@ -193,6 +279,7 @@ pub(super) mod fake { use super::*; use crate::errors::CometError; use parking_lot::Mutex; + use std::sync::atomic::{AtomicBool, Ordering::Relaxed}; /// Grants up to `limit` bytes in total and records every release. Panics if it is handed /// back more than it granted, which is the guarantee the pools must keep. @@ -201,7 +288,8 @@ pub(super) mod fake { limit: Mutex, held: Mutex, released: Mutex>, - fail: bool, + requests: Mutex>, + fail: AtomicBool, during_acquire: Mutex>>, } @@ -215,11 +303,16 @@ pub(super) mod fake { pub(in crate::execution::memory_pools) fn failing() -> Arc { Arc::new(Self { - fail: true, + fail: AtomicBool::new(true), ..Default::default() }) } + /// Makes every acquire from now on fail, or succeed again. + pub(in crate::execution::memory_pools) fn fail_acquires(&self, fail: bool) { + self.fail.store(fail, Relaxed); + } + pub(in crate::execution::memory_pools) fn set_limit(&self, limit: usize) { *self.limit.lock() = limit; } @@ -242,6 +335,11 @@ pub(super) mod fake { self.released.lock().clone() } + /// The size of every acquire, in the order they reached Spark. + pub(in crate::execution::memory_pools) fn requests(&self) -> Vec { + self.requests.lock().clone() + } + pub(in crate::execution::memory_pools) fn memory(self: &Arc) -> SparkMemory { SparkMemory::with_manager(Box::new(Arc::clone(self)), 0) } @@ -249,12 +347,13 @@ pub(super) mod fake { impl SparkMemoryManager for Arc { fn acquire(&self, size: usize) -> CometResult { + self.requests.lock().push(size); // Run before taking any other lock, so `f` can call back into this fake. let during = self.during_acquire.lock().take(); if let Some(f) = during { f(); } - if self.fail { + if self.fail.load(Relaxed) { return Err(CometError::Internal("jni".to_string())); } let limit = *self.limit.lock(); @@ -356,6 +455,193 @@ mod tests { assert_eq!(fake.held(), 0); } + /// Spark grants 0 of a 5 byte grow, so 5 bytes are owed and Spark then has room for more. + fn five_bytes_owed() -> (Arc, Arc) { + let fake = FakeSpark::with(0); + let spark = Arc::new(fake.memory()); + spark.acquire(5); + assert_eq!(spark.overcommit(), 5); + fake.set_limit(1_000); + (fake, spark) + } + + #[test] + fn concurrent_try_acquires_ask_spark_for_the_debt_once() { + let (fake, spark) = five_bytes_owed(); + // A second try_acquire runs while the first waits on Spark with the debt in its request. + let other = Arc::clone(&spark); + fake.during_next_acquire(move || assert_eq!(other.try_acquire(1).unwrap(), Ok(()))); + assert_eq!(spark.try_acquire(1).unwrap(), Ok(())); + assert_eq!(fake.requests(), vec![5, 6, 1]); + assert!(fake.released().is_empty(), "nothing was granted twice"); + assert_eq!(spark.overcommit(), 0); + assert_eq!(spark.claimed(), 0); + assert_eq!(fake.held(), 7); + } + + #[test] + fn refused_try_acquire_leaves_the_debt_for_the_next_one() { + let (fake, spark) = five_bytes_owed(); + fake.set_limit(3); + assert_eq!( + spark.try_acquire(1).unwrap(), + Err(Refusal { + overcommit: 5, + granted: 3 + }) + ); + assert_eq!(spark.claimed(), 0); + fake.set_limit(1_000); + assert_eq!(spark.try_acquire(1).unwrap(), Ok(())); + assert_eq!(fake.requests(), vec![5, 6, 6]); + assert_eq!(spark.overcommit(), 0); + assert_eq!(fake.held(), 6); + } + + #[test] + fn failed_try_acquire_leaves_the_debt_for_the_next_one() { + let (fake, spark) = five_bytes_owed(); + fake.fail_acquires(true); + assert!(spark.try_acquire(1).is_err()); + assert_eq!(spark.claimed(), 0); + fake.fail_acquires(false); + assert_eq!(spark.try_acquire(1).unwrap(), Ok(())); + assert_eq!(fake.requests(), vec![5, 6, 6]); + assert_eq!(spark.overcommit(), 0); + assert_eq!(fake.held(), 6); + } + + #[test] + fn panicking_try_acquire_leaves_the_debt_for_the_next_one() { + let (fake, spark) = five_bytes_owed(); + fake.during_next_acquire(|| panic!("injected Spark panic")); + let panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| spark.try_acquire(1))); + assert!(panic.is_err()); + assert_eq!(spark.claimed(), 0); + assert_eq!(spark.try_acquire(1).unwrap(), Ok(())); + assert_eq!(fake.requests(), vec![5, 6, 6]); + assert_eq!(spark.overcommit(), 0); + assert_eq!(fake.held(), 6); + } + + /// A release shrinks the debt under a claim of 5 bytes to 1, and a grow that Spark grants + /// none of then adds 3. The claim still covers all 4 bytes owed, so a try_acquire made + /// meanwhile claims none of them, and the claim repays all 4 when it settles. + #[test] + fn try_acquire_claims_no_debt_an_earlier_claim_still_covers() { + let (fake, spark) = five_bytes_owed(); + let (inner_fake, inner_spark) = (Arc::clone(&fake), Arc::clone(&spark)); + fake.during_next_acquire(move || { + inner_spark.release(4).unwrap(); + assert_eq!(inner_spark.overcommit(), 1); + assert_eq!(inner_spark.claimed(), 5); + inner_fake.set_limit(0); + inner_spark.acquire(3); + inner_fake.set_limit(1_000); + assert_eq!(inner_spark.overcommit(), 4); + assert_eq!(inner_spark.try_acquire(1).unwrap(), Ok(())); + assert_eq!(inner_spark.overcommit(), 4); + assert_eq!(inner_spark.claimed(), 5); + }); + assert_eq!(spark.try_acquire(1).unwrap(), Ok(())); + + // The second try_acquire asked for its own byte only. + assert_eq!(fake.requests(), vec![5, 6, 3, 1]); + // The claim repaid the 4 bytes owed, and Spark gets back the byte the release repaid. + assert_eq!(fake.released(), vec![1]); + assert_eq!(spark.overcommit(), 0); + assert_eq!(spark.claimed(), 0); + // 5 + 1 + 3 + 1 recorded and 4 released. + assert_eq!(fake.held() + spark.overcommit(), 6); + spark.release(6).unwrap(); + assert_eq!(fake.held(), 0); + } + + /// Two try_acquires are in flight, the first asking for 3 bytes owed and the second for the + /// 2 owed since. A release then repays 4 of the 5 bytes. Whichever settles first repays the + /// byte still owed, and Spark gets back exactly the bytes it granted that nothing owes. + #[test] + fn release_during_two_claims_hands_spark_back_only_the_excess() { + let fake = FakeSpark::with(10); + let spark = Arc::new(fake.memory()); + spark.acquire(13); + assert_eq!(spark.overcommit(), 3); + + let (inner_fake, inner_spark) = (Arc::clone(&fake), Arc::clone(&spark)); + fake.during_next_acquire(move || { + // 2 more bytes owed while the first claim waits on Spark. + inner_fake.set_limit(13); + inner_spark.acquire(5); + assert_eq!(inner_spark.overcommit(), 5); + inner_fake.set_limit(1_000); + // The second claim waits on Spark when the release lands. + let releaser = Arc::clone(&inner_spark); + inner_fake.during_next_acquire(move || releaser.release(4).unwrap()); + assert_eq!(inner_spark.try_acquire(1).unwrap(), Ok(())); + }); + assert_eq!(spark.try_acquire(1).unwrap(), Ok(())); + + assert_eq!(fake.requests(), vec![13, 4, 5, 3]); + // The second settled first and repaid the byte still owed: 1 of its 2 bytes and all 3 + // of the first's go back. + assert_eq!(fake.released(), vec![1, 3]); + assert_eq!(spark.overcommit(), 0); + assert_eq!(spark.claimed(), 0); + // 13 + 5 + 1 + 1 recorded and 4 released. + assert_eq!(fake.held(), 16); + spark.release(16).unwrap(); + assert_eq!(fake.held(), 0); + } + + /// Threads acquire, try_acquire and release at random while Spark's headroom moves. FakeSpark + /// panics if it is handed back more than it granted, and once the threads stop, Spark's + /// grant plus the overcommit must equal the bytes recorded. + #[test] + fn concurrent_callers_keep_the_ledger_balanced() { + use rand::{rngs::StdRng, RngExt, SeedableRng}; + use std::thread; + + let fake = FakeSpark::with(500); + let spark = Arc::new(fake.memory()); + let threads: Vec<_> = (0..4) + .map(|seed| { + let (fake, spark) = (Arc::clone(&fake), Arc::clone(&spark)); + thread::spawn(move || { + let mut rng = StdRng::seed_from_u64(seed); + let mut recorded = 0; + for _ in 0..5_000 { + let size = rng.random_range(1..100); + match rng.random_range(0..4) { + 0 => { + spark.acquire(size); + recorded += size; + } + 1 => { + if spark.try_acquire(size).unwrap().is_ok() { + recorded += size; + } + } + 2 if recorded > 0 => { + let size = rng.random_range(1..=recorded); + spark.release(size).unwrap(); + recorded -= size; + } + _ => fake.set_limit(rng.random_range(0..1_000)), + } + } + recorded + }) + }) + .collect(); + let recorded: usize = threads.into_iter().map(|t| t.join().unwrap()).sum(); + + assert_eq!(spark.claimed(), 0); + assert_eq!(fake.held() + spark.overcommit(), recorded); + spark.release(recorded).unwrap(); + assert_eq!(fake.held(), 0); + assert_eq!(spark.overcommit(), 0); + } + #[test] fn failed_acquire_is_all_overcommit() { let fake = FakeSpark::failing(); diff --git a/native/core/src/execution/memory_pools/unified_pool.rs b/native/core/src/execution/memory_pools/unified_pool.rs index fdb76bd4862..a9793d3e0a4 100644 --- a/native/core/src/execution/memory_pools/unified_pool.rs +++ b/native/core/src/execution/memory_pools/unified_pool.rs @@ -119,10 +119,11 @@ impl MemoryPool for CometUnifiedMemoryPool { // A partial grant is handed back and refused, which triggers spilling in the caller. if let Err(refusal) = self.spark.try_acquire(additional)? { return Err(resources_datafusion_err!( - "Task {} failed to acquire {} bytes plus {} bytes overcommitted, only got {}. Reserved: {}", + "Task {} failed to acquire {} bytes plus {} of {} bytes overcommitted, only got {}. Reserved: {}", self.spark.task_attempt_id(), additional, refusal.overcommit, + self.spark.overcommit(), refusal.granted, self.reserved() )); @@ -178,6 +179,26 @@ mod tests { assert_eq!(fake.held(), 0); } + /// A try_grow refused while another one in flight is asking Spark for the debt reports the + /// overcommit outstanding, not only the part it asked for itself. + #[test] + fn try_grow_refusal_reports_debt_another_try_grow_is_asking_for() { + let fake = FakeSpark::with(100); + let pool: Arc = Arc::new(CometUnifiedMemoryPool::with_spark(fake.memory())); + let first = MemoryConsumer::new("first").register(&pool); + let second = MemoryConsumer::new("second").register(&pool); + // Spark grants 100 of the 150 bytes, so 50 are owed and Spark is full. + first.grow(150); + + let (err_tx, err_rx) = std::sync::mpsc::channel(); + fake.during_next_acquire(move || { + let _ = err_tx.send(second.try_grow(1).unwrap_err().to_string()); + }); + assert!(first.try_grow(5).is_err()); + let err = err_rx.recv().unwrap(); + assert!(err.contains("plus 0 of 50 bytes overcommitted"), "{err}"); + } + #[test] fn concurrent_consumers_hand_spark_back_exactly_what_it_granted() { use rand::{rngs::StdRng, RngExt, SeedableRng}; diff --git a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java index 1c76627c5b5..50aecdddef2 100644 --- a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java +++ b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java @@ -50,6 +50,8 @@ public class CometTaskMemoryManager { public final TaskMemoryManager internal; private final NativeMemoryConsumer nativeMemoryConsumer; + + /** Bytes Comet's memory pools hold from Spark through this manager, see {@link #getUsed}. */ private final AtomicLong used = new AtomicLong(); public CometTaskMemoryManager(long id, long taskAttemptId) { @@ -120,6 +122,10 @@ public void releaseMemory(long size) { if (logger.isTraceEnabled()) { logger.trace("Task {} released {} bytes", taskAttemptId, size); } + // Spark is called first and `used` moves only once it returns, so a release that throws has + // moved nothing on Comet's side: the native pool still counts those bytes as held and so + // does `used`. + internal.releaseExecutionMemory(size, nativeMemoryConsumer); long newUsed = used.addAndGet(-size); if (newUsed < 0) { logger.error( @@ -128,9 +134,12 @@ public void releaseMemory(long size) { newUsed, size); } - internal.releaseExecutionMemory(size, nativeMemoryConsumer); } + /** + * Bytes Comet's memory pools hold from Spark through this manager. A non-zero value once the + * task's last native plan has been released is a reservation that was never freed. + */ public long getUsed() { return used.get(); } diff --git a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala index 82c5731da7a..a3e17d7a3fe 100644 --- a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala +++ b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala @@ -19,8 +19,8 @@ package org.apache.spark -import java.util.concurrent.TimeUnit -import java.util.concurrent.atomic.AtomicReference +import java.util.concurrent.{CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.{AtomicBoolean, AtomicLong, AtomicReference} import org.apache.logging.log4j.Level import org.apache.logging.log4j.core.LogEvent @@ -75,6 +75,38 @@ class CometTaskMemoryManagerSuite extends TaskMemoryTestUtils { } } + test("a release that Spark rejects leaves the usage unchanged") { + val memoryManager = new TestMemoryManager(new SparkConf()) + memoryManager.limit(1024) + val rejectRelease = new AtomicBoolean(false) + val taskMemoryManager = new TaskMemoryManager(memoryManager, 0L) { + override def releaseExecutionMemory(size: Long, consumer: MemoryConsumer): Unit = { + if (rejectRelease.get) { + throw new IllegalStateException("release rejected") + } + super.releaseExecutionMemory(size, consumer) + } + } + + withTaskContext(taskMemoryManager) { _ => + val manager = new CometTaskMemoryManager(1L, 0L) + val consumer = nativeMemoryConsumer(manager) + assert(manager.acquireMemory(128L) == 128L) + + rejectRelease.set(true) + intercept[IllegalStateException](manager.releaseMemory(64L)) + rejectRelease.set(false) + // Native code gets the error and still counts every byte as held, and so must the usage. + assert(manager.getUsed == 128L) + assert(consumer.getUsed == 128L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 128L) + + manager.releaseMemory(128L) + assert(manager.getUsed == 0L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + } + } + test("managers in the same task retain separate memory accounting") { withTaskMemoryManager { taskMemoryManager => val first = new CometTaskMemoryManager(1L, 0L) @@ -134,6 +166,78 @@ class CometTaskMemoryManagerSuite extends TaskMemoryTestUtils { } } + test("a short grant is handed back while another acquire of the task waits in Spark") { + // Another task holds 82 bytes and a third holds 1. + val memoryManager = offHeapMemoryManager() + val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L)) + val thirdTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 2L)) + assert(otherTask.acquireMemory(82L) == 82L) + assert(thirdTask.acquireMemory(1L) == 1L) + + val shortGrant = new CountDownLatch(1) + val secondParked = new CountDownLatch(1) + val taskMemoryManager = new TaskMemoryManager(memoryManager, 0L) { + override def acquireExecutionMemory(required: Long, consumer: MemoryConsumer): Long = { + val got = super.acquireExecutionMemory(required, consumer) + // Out of Spark's monitor, the short grant waits until the second acquire has parked. + if (got < required && shortGrant.getCount > 0) { + shortGrant.countDown() + secondParked.await(2 * TimeoutSeconds, TimeUnit.SECONDS) + } + got + } + } + + withTaskContext(taskMemoryManager) { _ => + val manager = new CometTaskMemoryManager(1L, 0L) + // Hold one byte so that handing back the short grant never empties the task's balance, + // which would drop its entry from Spark's pool under the waiting acquire. + assert(manager.acquireMemory(1L) == 1L) + + // Three active tasks and 16 bytes free: a 30 byte request is short granted 16 bytes, which + // native code then hands back. + val first = new NativeThread( + "the first acquire", + () => { + val granted = manager.acquireMemory(30L) + manager.releaseMemory(granted) + granted + }) + val second = new NativeThread("the second acquire", () => manager.acquireMemory(10L)) + + try { + // At DEBUG the short grant is logged, which must not wait on the task's monitor that the + // parked second acquire holds. + val messages = logEvents(Level.DEBUG) { + first.start() + assert(shortGrant.await(TimeoutSeconds, TimeUnit.SECONDS), s"no short grant: $first") + // The third task leaves. With two active tasks this task's minimum share is 25 bytes + // and 1 byte is free, so a 10 byte request waits inside Spark holding the task's + // monitor. + thirdTask.freeMemory(1L) + second.start() + withClue(second)(awaitWaitingInSpark(second.thread)) + secondParked.countDown() + + // Handing back the short grant is what lets the second acquire through. + assert(second.result(first) == 10L) + assert(first.result(second) == 16L) + }.map(_.getMessage.getFormattedMessage) + assert( + messages.exists(_.contains("requested 30 bytes but only received 16 bytes")), + messages.mkString("\n")) + } finally { + secondParked.countDown() + // Free the other task's memory so that neither thread outlives a failed test. + otherTask.freeMemory(otherTask.getUsed) + Seq(first, second).foreach(_.join()) + } + manager.releaseMemory(10L) + manager.releaseMemory(1L) + assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L) + } + } + test("an acquire waiting in Spark survives a release that empties the task's balance") { checkRequestWhileBalanceEmpties(loggers) { taskMemoryManager => val manager = new CometTaskMemoryManager(1L, 0L) @@ -255,4 +359,28 @@ class CometTaskMemoryManagerSuite extends TaskMemoryTestUtils { field.setAccessible(true) field.get(manager).asInstanceOf[MemoryConsumer] } + + /** Runs `body` on a daemon thread, as native code would, keeping its result or failure. */ + private class NativeThread(name: String, body: () => Long) { + private val value = new AtomicLong(-1L) + private val failure = new AtomicReference[Throwable]() + val thread: Thread = new Thread(() => + try value.set(body()) + catch { case t: Throwable => failure.set(t) }) + thread.setDaemon(true) + + def start(): Unit = thread.start() + + def join(): Unit = thread.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds)) + + /** Waits for `body` to return and gives its result, adding `clue` to a failure. */ + def result(clue: => Any = ""): Long = { + join() + assert(!thread.isAlive && failure.get == null, s"$this. $clue") + value.get + } + + override def toString: String = + Option(failure.get).fold(s"$name is ${thread.getState}")(t => s"$name failed: $t") + } }