dwsmith1983 commented on code in PR #5613:
URL: https://github.com/apache/datafusion-comet/pull/5613#discussion_r4218458410
##########
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:
The teardown wait and the `TeardownComplete` guard are gone along with the
anchor, so `task_shared.rs` matches main again, `TaskSharedMemoryPool::drop`
compares pointers as before, and 155e74060 puts the skill back to main's
wording for that bullet.
--
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]
