andygrove commented on code in PR #5613:
URL: https://github.com/apache/datafusion-comet/pull/5613#discussion_r4178374734
##########
native/core/src/execution/memory_pools/task_shared.rs:
##########
@@ -81,37 +113,33 @@ impl MemoryPool for TaskSharedMemoryPool {
}
}
-impl Drop for TaskSharedMemoryPool {
- fn drop(&mut self) {
- if let Entry::Occupied(entry) =
TASK_SHARED_MEMORY_POOLS.lock().entry(self.task_attempt_id)
- {
- // An acquire racing with this drop can replace our expired `Weak`
before we obtain the
- // lock. Do not let the old pool remove that replacement's entry.
- if std::ptr::eq(entry.get().as_ptr(), self) {
- entry.remove();
- }
- }
- }
-}
-
/// Returns the memory pool shared by every native plan in `task_attempt_id`,
creating it with
/// `create` if no live pool exists for the task. The returned `Arc` is the
RAII handle: the pool
-/// stays registered until the last reference to it drops.
+/// stays registered until the last reference to it drops. `create` runs under
the shared registry
+/// lock, so it must make no JVM call and must not block. `create` must return
a pool nothing else
+/// holds a reference to, because the wait below relies on dropping the
wrapper dropping the pool
+/// itself.
+///
+/// While the task's previous pool is still dropping, this waits for that drop
to finish before it
+/// creates the replacement. The fair pool hands its anchor byte back to Spark
from its drop, and a
+/// replacement created before that release lands holds nothing from Spark
yet, so its first
+/// acquire could park inside Spark and wake to a task entry the release has
removed.
pub(crate) fn acquire_task_shared_pool(
task_attempt_id: i64,
create: impl FnOnce() -> Arc<dyn MemoryPool>,
) -> Arc<dyn MemoryPool> {
let mut memory_pool_map = TASK_SHARED_MEMORY_POOLS.lock();
- if let Some(memory_pool) = memory_pool_map
- .get(&task_attempt_id)
- .and_then(Weak::upgrade)
- {
- return memory_pool;
+ loop {
+ match memory_pool_map.get(&task_attempt_id).map(Weak::upgrade) {
+ Some(Some(memory_pool)) => return memory_pool,
+ Some(None) => TEARDOWN_COMPLETE.wait(&mut memory_pool_map),
Review Comment:
The description still says the registry is untouched and that
`task_shared.rs` and `mod.rs` only differ from main by doc sentences. Since
`ac0a691a0`, `acquire_task_shared_pool` waits on a process-wide `Condvar` while
the task's previous pool is still dropping, and the `TeardownComplete` guard
replaces `TaskSharedMemoryPool::drop`. Could you add a paragraph on that, and
on why the replacement has to wait for the anchor release? It is a new blocking
wait in `createPlan`, so readers of the PR should know it is there.
##########
native/core/src/execution/memory_pools/task_shared.rs:
##########
@@ -19,22 +19,54 @@ use datafusion::execution::memory_pool::{
MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation,
};
use once_cell::sync::Lazy;
-use parking_lot::Mutex;
+use parking_lot::{Condvar, Mutex};
use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::fmt;
use std::sync::{Arc, Weak};
/// The memory pools for active task attempts. Weak references let the pool's
normal `Arc`
-/// ownership determine its lifetime, and each pool removes its entry when the
last reference drops.
+/// ownership determine its lifetime. An entry whose `Weak` has expired
belongs to a pool that is
+/// still dropping; it is removed once that drop has finished.
static TASK_SHARED_MEMORY_POOLS: Lazy<Mutex<HashMap<i64,
Weak<TaskSharedMemoryPool>>>> =
Lazy::new(|| Mutex::new(HashMap::new()));
+/// Notified under `TASK_SHARED_MEMORY_POOLS` each time a pool has finished
dropping.
+static TEARDOWN_COMPLETE: Condvar = Condvar::new();
+
/// A transparent `MemoryPool` wrapper whose lifetime also controls its
registry entry.
#[derive(Debug)]
struct TaskSharedMemoryPool {
- task_attempt_id: i64,
inner: Arc<dyn MemoryPool>,
+ /// Declared after `inner` on purpose. Fields drop in declaration order,
so the inner pool,
+ /// including the anchor byte a `CometFairMemoryPool` hands back to Spark
from its drop, is
+ /// gone before this guard removes the registry entry and wakes the
acquires waiting for it.
+ _teardown: TeardownComplete,
+}
+
+/// Removes a pool's registry entry once the pool itself has dropped, and
wakes any acquire
+/// waiting for that.
+#[derive(Debug)]
+struct TeardownComplete {
+ task_attempt_id: i64,
+}
+
+impl Drop for TeardownComplete {
+ fn drop(&mut self) {
+ let mut memory_pool_map = TASK_SHARED_MEMORY_POOLS.lock();
+ if let Entry::Occupied(entry) =
memory_pool_map.entry(self.task_attempt_id) {
+ // In production only this guard removes entries, and
`acquire_task_shared_pool`
+ // creates a pool only when no entry exists, so the entry found
here is this pool's
+ // own, expired. The liveness check is a safeguard for an entry
replaced out of band,
+ // as the test `an_old_pool_does_not_remove_its_replacement` does.
It reads
+ // `strong_count` rather than calling `upgrade` because dropping
an upgraded `Arc`
+ // under the lock could be the last reference.
+ if entry.get().strong_count() == 0 {
Review Comment:
`.ai/skills/review-comet-memory-pr/SKILL.md` still says that
`TaskSharedMemoryPool::drop` compares pointers so an old pool never removes its
replacement's entry, and asks reviewers not to let that check be simplified
away. This PR removes that `Drop` impl. The race can't happen anymore because
`acquire_task_shared_pool` waits for the dying pool, and `TeardownComplete`
only removes an expired entry. Could you update that bullet here so it
describes the wait and the guard? Otherwise the next review that loads the
skill will flag this design as a regression.
--
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]
