comments

Author: xanderbailey

datafusion/physical-plan/src/spill/spill_pool.rs

@@ -91,35 +91,33 @@ 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.
+/// Cloning increments the count. [`WriterCount::decrement`] atomically
+/// decrements the count and reports whether the caller was the last clone.
+/// This combined operation avoids TOCTOU races between checking the count
+/// and decrementing it.
 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
+    /// Decrements the count and returns `true` if this was the last clone.
+    ///
+    /// This is a single atomic operation, so concurrent drops cannot both
+    /// observe themselves as "last".
+    fn decrement(&self) -> bool {
+        self.0.fetch_sub(1, Ordering::SeqCst) == 1
     }
 }
 
 impl Clone for WriterCount {
     fn clone(&self) -> Self {
-        self.0.fetch_add(1, Ordering::Relaxed);
+        self.0.fetch_add(1, Ordering::SeqCst);
         Self(Arc::clone(&self.0))
     }
 }
 
-impl Drop for WriterCount {
-    fn drop(&mut self) {
-        self.0.fetch_sub(1, Ordering::Release);
-    }
-}
-
 /// Writer for a spill pool. Provides coordinated write access with FIFO semantics.
 ///
 /// Created by [`channel`]. See that function for architecture diagrams and usage examples.
@@ -266,10 +264,9 @@ impl SpillPoolWriter {
 
 impl Drop for SpillPoolWriter {
     fn drop(&mut self) {
-        if !self.writer_count.is_last() {
+        if !self.writer_count.decrement() {
             // Other writer clones are still active; do not finalize or
-            // signal EOF to readers.  `self.writer_count` is decremented
-            // automatically when this `Drop` returns.
+            // signal EOF to readers.
             return;
         }