dwsmith1983 commented on code in PR #5613:
URL: https://github.com/apache/datafusion-comet/pull/5613#discussion_r4178506016
##########
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:
> Could you update that bullet here so it describes the wait and the guard?
Done in a9ef739db. The bullet now says the entry stays with an expired
`Weak` until the pool has finished dropping, `acquire_task_shared_pool` waits
on `TEARDOWN_COMPLETE` while it does, and only the `TeardownComplete` guard,
declared after the inner pool, removes an expired entry. It asks reviewers not
to reorder those fields, remove the entry early, create a replacement without
waiting, or let `create` return a pool something else still references.
--
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]
