alamb commented on code in PR #20672:
URL: https://github.com/apache/datafusion/pull/20672#discussion_r2886560181
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -93,6 +89,37 @@ impl SpillPoolShared {
}
}
+/// Tracks the number of live [`SpillPoolWriter`] clones.
+///
+/// Cloning increments the count; dropping decrements it.
+/// [`WriterCount::is_last`] returns `true` when called from the final clone,
+/// which the writer uses to decide whether to finalize the spill pool.
+struct WriterCount(Arc<AtomicUsize>);
+
+impl WriterCount {
+ fn new() -> Self {
+ Self(Arc::new(AtomicUsize::new(1)))
+ }
+
+ /// Returns `true` if this is the only remaining clone.
+ fn is_last(&self) -> bool {
+ self.0.load(Ordering::Acquire) == 1
+ }
+}
+
+impl Clone for WriterCount {
+ fn clone(&self) -> Self {
+ self.0.fetch_add(1, Ordering::Relaxed);
Review Comment:
I think using Ordering::SeqCst or Ordering::Acquire might be safer here. It
probably is unlikely to cause problems in practice, but I think it would be
important that any wrtiter increment is seen before the `is_last` check is done
https://doc.rust-lang.org/std/sync/atomic/enum.Ordering.html
##########
datafusion/physical-plan/src/spill/spill_pool.rs:
##########
@@ -247,17 +266,15 @@ impl SpillPoolWriter {
impl Drop for SpillPoolWriter {
fn drop(&mut self) {
- let mut shared = self.shared.lock();
-
- shared.active_writer_count -= 1;
- let is_last_writer = shared.active_writer_count == 0;
-
- if !is_last_writer {
+ if !self.writer_count.is_last() {
Review Comment:
I think there a race condition now -- what prevents a new writer from being
added after `is_last` returns false but before the lock is taken below?
--
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]
