dwsmith1983 commented on code in PR #5613:
URL: https://github.com/apache/datafusion-comet/pull/5613#discussion_r4178509429
##########
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:
> Could you add a paragraph on that, and on why the replacement has to wait
for the anchor release?
Added under "What changes are included". It covers why the replacement waits
(a pool created before the anchor release lands holds nothing from Spark, so
its first acquire can park and wake to a removed task entry), the
`TeardownComplete` guard replacing `TaskSharedMemoryPool::drop`, and that
`createPlan` can now block, only while the task's pool is mid-drop and until
that drop returns. The tests paragraph now lists the replacement tests as well.
--
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]
