mbutrovich commented on code in PR #5613:
URL: https://github.com/apache/datafusion-comet/pull/5613#discussion_r4360827020
##########
native/core/src/execution/memory_pools/fair_pool.rs:
##########
@@ -202,9 +278,354 @@ 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;
+
+ const THIS_TASK: i64 = 0;
+ const OTHER_TASK: i64 = 1;
+ const THIRD_TASK: i64 = 2;
+
+ /// Real Spark parks a starved acquire forever; the stub panics after this
so a deadlock
+ /// fails the test instead of hanging it.
+ const WAIT_TIMEOUT: Duration = Duration::from_secs(5);
+ const TEST_TIMEOUT: Duration = Duration::from_secs(20);
+
+ /// While armed, holds a bridge call until the test opens it, to
interleave another thread.
+ struct Gate {
+ armed: AtomicBool,
+ entered: (Sender<()>, Mutex<Receiver<()>>),
+ open: (Sender<()>, Mutex<Receiver<()>>),
+ }
+
+ 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);
+ }
+
+ 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<i64, i64>,
+ }
+
+ impl SparkPool {
+ fn memory_free(&self) -> i64 {
+ self.pool_size - self.memory_for_task.values().sum::<i64>()
+ }
+ }
+
+ /// Models Spark 4.1.3's `ExecutionMemoryPool` acquire and release for
this task: the entry
+ /// lifecycle, the 1/N and 1/(2N) shares and the wait loop. Re-adding the
entry in the wait
+ /// loop stands in for the `CometTaskMemoryManager` retry, which these
tests do not cover.
+ struct StubTaskMemory {
Review Comment:
Should the stub model the task monitor that `acquireExecutionMemory` holds,
for example a per-task mutex held across `acquire`?
At the moment the stub lets two acquires of one task be inside Spark at
once. The JVM can't reach that state. The tests that rely on it check behavior
that can't happen in production:
- `second_acquire_does_not_queue_behind_a_parked_one`
- `concurrent_try_grows_ask_spark_for_the_overcommit_once`
- `short_grant_rollback_landing_under_a_parked_sibling_wakes_it`
- `try_grow_refusal_reports_debt_another_try_grow_is_asking_for`, whose
`during_next_acquire` callback calls `try_grow` from inside the first acquire
on the same thread
Several earlier rounds found that the stub differed from real Spark. A stub
that also serializes acquires would keep the tests honest about which
interleavings are reachable.
##########
native/core/src/execution/memory_pools/fair_pool.rs:
##########
@@ -152,45 +196,77 @@ 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();
+ // Check and charge under one lock so concurrent grows cannot
jointly exceed a share
+ // or pool_size.
Review Comment:
What do you think about splitting the pool's lock in two instead, the way
Spark's `TaskMemoryManager` does?
`acquireExecutionMemory` holds the task's monitor for the whole call,
including while `ExecutionMemoryPool.acquireMemory` is parked in `lock.wait()`
([TaskMemoryManager.java#L205](https://github.com/apache/spark/blob/4bc916c99a46cc525d6b05458bbadedcdc43a0f7/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java#L205),
[ExecutionMemoryPool.scala#L139-L142](https://github.com/apache/spark/blob/4bc916c99a46cc525d6b05458bbadedcdc43a0f7/core/src/main/scala/org/apache/spark/memory/ExecutionMemoryPool.scala#L139-L142)).
`releaseExecutionMemory` never takes that monitor
([TaskMemoryManager.java#L438-L452](https://github.com/apache/spark/blob/4bc916c99a46cc525d6b05458bbadedcdc43a0f7/core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java#L438-L452)).
So in the JVM a task's acquires already run one at a time, and only releases
run alongside them. @andygrove pointed this out in an earlier round, and the
doc comment on `second_acquire_does_not_queue_behind_a_parked_one` s
ays so too.
The pool could copy that split:
- An `acquire: Mutex<()>` that `grow` and `try_grow` hold across the Spark
call. It costs no concurrency, because Spark would queue the second acquire on
its monitor anyway.
- The existing `state` mutex, held only for arithmetic. `shrink` and
`reserved()` take only this one, so neither waits behind a parked acquire.
Against `main`, the change looks like this:
```rust
// try_grow
let _acquire = self.acquire.lock();
let state = self.state.lock();
// ... the two limit checks, unchanged ...
drop(state);
if let Err(refusal) = self.spark.try_acquire(additional)? { /* refuse, as on
main */ }
let mut state = self.state.lock();
state.used += additional;
*state.consumer_used(reservation) += additional;
// shrink: check under `state`, release to Spark with no lock held, then
subtract under `state`
```
With one acquire in flight, a limit check stays valid until its charge
lands, because a concurrent shrink can only lower the totals. That removes the
provisional charge, `settle`, `try_acquire_leaving_a_short_grant` and the claim
ledger from the fair pool. It also removes most of the "two windows" section
the guide now needs to describe them.
Each round of this review has found a new race in the state that concurrent
acquires leave in flight. I'd rather not keep that state when Spark doesn't let
it happen. I know this reopens the charge-then-settle shape that earlier rounds
accepted. I'd still rather change it before merge, because every later reader
of this pool would have to reason about that in-flight state.
The other fixes in this PR address bugs that exist on `main`: the unified
pool's debt race, keeping a short grant charged when its rollback release
fails, and the `releaseMemory` reorder in `CometTaskMemoryManager`. Could each
go in its own PR, so this one does what its title says?
##########
native/core/src/execution/memory_pools/spark_memory.rs:
##########
@@ -58,34 +56,81 @@ 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 release while a claim is parked in Spark leaves that request too large,
so it may come back
+/// short and be refused. Handing the claim to Spark on release instead 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<dyn SparkMemoryManager>,
- overcommit: AtomicUsize,
+ /// Held only for arithmetic, never across a call to Spark.
+ ledger: Mutex<Ledger>,
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,
+}
+
Review Comment:
Could the ledger and `Claim` move to their own PR, with an issue for the
unified pool?
On `main`, `CometUnifiedMemoryPool` has no lock, and `try_acquire` reads the
debt with a plain load
([spark_memory.rs#L111-L112](https://github.com/apache/datafusion-comet/blob/6a6d71c52f3452ac92ebe1c30c92c8515508f59e/native/core/src/execution/memory_pools/spark_memory.rs#L111-L112)).
So two concurrent unified `try_grow` calls can both ask Spark for the same
debt today. That fix stands on its own and deserves its own review. If the fair
pool serializes its acquires, the fair pool no longer needs the ledger.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
