reclaimer state

Author: JanKaul

datafusion/execution/src/memory_pool/mod.rs

@@ -39,7 +39,7 @@ pub use datafusion_common::{
     human_readable_count, human_readable_duration, human_readable_size, units,
 };
 pub use pool::*;
-pub use reclaimer::MemoryReclaimer;
+pub use reclaimer::{MemoryReclaimer, reclaimer_state};
 
 /// Tracks and potentially limits memory use across operators during execution.
 ///

datafusion/execution/src/memory_pool/pool.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use crate::memory_pool::reclaimer::reclaimer_state;
 use crate::memory_pool::{
     MemoryConsumer, MemoryLimit, MemoryPool, MemoryReclaimer, MemoryReservation,
     human_readable_size,
@@ -29,7 +30,7 @@ use std::pin::Pin;
 use std::sync::Arc;
 use std::{
     num::NonZeroUsize,
-    sync::atomic::{AtomicBool, AtomicUsize, Ordering},
+    sync::atomic::{AtomicU8, AtomicUsize, Ordering},
 };
 
 /// A [`MemoryPool`] that enforces no limit
@@ -329,34 +330,49 @@ struct TrackedConsumer {
     reserved: AtomicUsize,
     peak: AtomicUsize,
     reclaimer: Option<Arc<dyn MemoryReclaimer>>,
-    /// Set to `true` while a reclaim is in-flight on this consumer, so
-    /// concurrent `try_grow_async` callers skip this victim instead of
-    /// double-reclaiming. Reset via [`ReclaimInFlightGuard`] on drop
-    /// (cancel-safe).
-    reclaim_in_flight: Arc<AtomicBool>,
+    /// Tri-state eligibility flag for [`reclaimer`], encoded per
+    /// [`reclaimer_state`]. The pool flips `AVAILABLE` ↔ `IN_FLIGHT`
+    /// for dedup; the reclaimer's owner may sticky-set `DISABLED` once
+    /// it can no longer free memory. Shared `Arc` so the reclaimer
+    /// side and the pool see the same cell. `None` reclaimer ⇒ flag
+    /// is unused but still allocated.
+    reclaimer_state: Arc<AtomicU8>,
 }
 
-/// RAII guard that clears [`TrackedConsumer::reclaim_in_flight`] on drop.
-struct ReclaimInFlightGuard {
-    flag: Arc<AtomicBool>,
+/// RAII guard for the [`IN_FLIGHT`] slot of a [`TrackedConsumer`]'s
+/// `reclaimer_state` flag. `Drop` only restores `AVAILABLE` if the
+/// state is still `IN_FLIGHT` — leaves a sticky `DISABLED` alone.
+///
+/// [`IN_FLIGHT`]: reclaimer_state::IN_FLIGHT
+struct ReclaimerStateGuard {
+    flag: Arc<AtomicU8>,
 }
 
-impl Drop for ReclaimInFlightGuard {
+impl Drop for ReclaimerStateGuard {
     fn drop(&mut self) {
-        self.flag.store(false, Ordering::Release);
+        let _ = self.flag.compare_exchange(
+            reclaimer_state::IN_FLIGHT,
+            reclaimer_state::AVAILABLE,
+            Ordering::AcqRel,
+            Ordering::Relaxed,
+        );
     }
 }
 
-impl ReclaimInFlightGuard {
-    /// Atomically claims the in-flight slot, returning `Some(guard)` on
-    /// success or `None` if another reclaim is already running on this
-    /// consumer.
-    fn try_acquire(flag: &Arc<AtomicBool>) -> Option<Self> {
-        flag.compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed)
-            .ok()
-            .map(|_| Self {
-                flag: Arc::clone(flag),
-            })
+impl ReclaimerStateGuard {
+    /// Try to transition the flag from `AVAILABLE` to `IN_FLIGHT`.
+    /// Fails on contention or on a sticky `DISABLED`.
+    fn try_acquire(flag: &Arc<AtomicU8>) -> Option<Self> {
+        flag.compare_exchange(
+            reclaimer_state::AVAILABLE,
+            reclaimer_state::IN_FLIGHT,
+            Ordering::AcqRel,
+            Ordering::Relaxed,
+        )
+        .ok()
+        .map(|_| Self {
+            flag: Arc::clone(flag),
+        })
     }
 }
 
@@ -598,6 +614,16 @@ impl<I: MemoryPool> MemoryPool for TrackConsumersPool<I> {
     fn register(&self, consumer: &MemoryConsumer) {
         self.inner.register(consumer);
 
+        let reclaimer = consumer.reclaimer().cloned();
+        // Reuse the reclaimer's own flag when it provides one — that
+        // way the reclaimer side can sticky-set `DISABLED` and the
+        // pool sees it on the next filter pass. Otherwise allocate a
+        // fresh `AVAILABLE` flag for in-flight dedup only.
+        let state = reclaimer
+            .as_ref()
+            .and_then(|r| r.reclaimer_state())
+            .unwrap_or_else(|| Arc::new(AtomicU8::new(reclaimer_state::AVAILABLE)));
+
         let mut guard = self.tracked_consumers.write();
         let existing = guard.insert(
             consumer.id(),
@@ -606,8 +632,8 @@ impl<I: MemoryPool> MemoryPool for TrackConsumersPool<I> {
                 can_spill: consumer.can_spill(),
                 reserved: Default::default(),
                 peak: Default::default(),
-                reclaimer: consumer.reclaimer().cloned(),
-                reclaim_in_flight: Arc::new(AtomicBool::new(false)),
+                reclaimer,
+                reclaimer_state: state,
             },
         );
 
@@ -690,11 +716,10 @@ impl<I: MemoryPool> MemoryPool for TrackConsumersPool<I> {
             // mutual-reclaim cycle (A targets B while B targets A) — at
             // most one side of any pair can hold strictly more memory,
             // so the other side has no candidates and surfaces an error
-            // for the caller's self-spill fallback. Skip zero-byte
-            // consumers and drop the read guard before awaiting any
-            // reclaim. Clone each victim's `reclaim_in_flight` flag so
-            // concurrent `try_grow_async` callers dedup against in-flight
-            // reclaims on the same victim.
+            // for the caller's self-spill fallback. Filter out anyone
+            // whose `reclaimer_state` flag is not `AVAILABLE` (in-flight or
+            // sticky-disabled). Drop the read guard before awaiting any
+            // reclaim.
             let requestor_id = reservation.consumer().id();
             let requestor_reserved = {
                 let guard = self.tracked_consumers.read();
@@ -703,7 +728,7 @@ impl<I: MemoryPool> MemoryPool for TrackConsumersPool<I> {
                     .map(|tc| tc.reserved())
                     .unwrap_or(0)
             };
-            let mut candidates: Vec<(usize, Arc<dyn MemoryReclaimer>, Arc<AtomicBool>)> = {
+            let mut candidates: Vec<(usize, Arc<dyn MemoryReclaimer>, Arc<AtomicU8>)> = {
                 let guard = self.tracked_consumers.read();
                 guard
                     .iter()
@@ -712,10 +737,15 @@ impl<I: MemoryPool> MemoryPool for TrackConsumersPool<I> {
                         if *cid == requestor_id || tc.reserved() <= requestor_reserved {
                             return None;
                         }
+                        if tc.reclaimer_state.load(Ordering::Acquire)
+                            != reclaimer_state::AVAILABLE
+                        {
+                            return None;
+                        }
                         Some((
                             tc.reserved(),
                             Arc::clone(reclaimer),
-                            Arc::clone(&tc.reclaim_in_flight),
+                            Arc::clone(&tc.reclaimer_state),
                         ))
                     })
                     .collect()
@@ -728,15 +758,15 @@ impl<I: MemoryPool> MemoryPool for TrackConsumersPool<I> {
             candidates.truncate(self.reclaim_candidate_limit.get());
 
             // For each candidate: try to claim its in-flight slot
-            // (skip on contention so we work on a different victim
-            // rather than serializing behind a sibling's reclaim);
-            // re-check `try_grow` before reclaiming in case a sibling
-            // already freed enough; reclaim; retry `try_grow`. The
-            // retry path goes through `self.try_grow`, which already
-            // updates the tracked consumer's atomic reservation — no
-            // manual accounting needed here.
-            for (_, reclaimer, in_flight) in candidates {
-                let _g = match ReclaimInFlightGuard::try_acquire(&in_flight) {
+            // (skip on contention or sticky-disabled so we work on a
+            // different victim rather than serializing behind a
+            // sibling's reclaim); re-check `try_grow` before reclaiming
+            // in case a sibling already freed enough; reclaim; retry
+            // `try_grow`. The retry path goes through `self.try_grow`,
+            // which already updates the tracked consumer's atomic
+            // reservation — no manual accounting needed here.
+            for (_, reclaimer, flag) in candidates {
+                let _g = match ReclaimerStateGuard::try_acquire(&flag) {
                     Some(g) => g,
                     None => continue,
                 };

datafusion/execution/src/memory_pool/reclaimer.rs

@@ -22,6 +22,22 @@
 
 use datafusion_common::Result;
 use std::fmt::Debug;
+use std::sync::Arc;
+use std::sync::atomic::AtomicU8;
+
+/// Encoded values stored in the [`reclaimer_state`] tri-state.
+///
+/// [`reclaimer_state`]: MemoryReclaimer::reclaimer_state
+pub mod reclaimer_state {
+    /// Reclaimer is idle and may be selected as a victim.
+    pub const AVAILABLE: u8 = 0;
+    /// A pool task is currently driving `reclaim` on this reclaimer.
+    pub const IN_FLIGHT: u8 = 1;
+    /// Reclaimer has been retired (e.g. operator entered a phase where
+    /// it can no longer free memory). Sticky — never returns to
+    /// `AVAILABLE`.
+    pub const DISABLED: u8 = 2;
+}
 
 /// Hook attached to a [`MemoryConsumer`] via
 /// [`MemoryConsumer::with_reclaimer`]. On
@@ -58,4 +74,15 @@ pub trait MemoryReclaimer: Send + Sync + Debug {
     fn priority(&self) -> i32 {
         0
     }
+
+    /// Optional shared tri-state flag controlling whether the pool
+    /// currently considers this reclaimer eligible. Values are defined
+    /// in [`reclaimer_state`]. Returning `Some(arc)` lets the
+    /// reclaimer's owner flip itself to `DISABLED` once it can no
+    /// longer free memory (e.g., on entering a merge phase), which
+    /// the pool observes immediately. Returning `None` lets the pool
+    /// allocate its own private flag — used only for in-flight dedup.
+    fn reclaimer_state(&self) -> Option<Arc<AtomicU8>> {
+        None
+    }
 }

datafusion/physical-plan/src/sorts/sort.rs

@@ -66,7 +66,7 @@ use datafusion_common::{
 };
 use datafusion_execution::TaskContext;
 use datafusion_execution::memory_pool::{
-    MemoryConsumer, MemoryReclaimer, MemoryReservation,
+    MemoryConsumer, MemoryReclaimer, MemoryReservation, reclaimer_state,
 };
 use datafusion_execution::runtime_env::RuntimeEnv;
 use datafusion_physical_expr::LexOrdering;
@@ -82,6 +82,10 @@ use log::{debug, trace};
 #[derive(Debug)]
 struct ExternalSorterReclaimer {
     tx: tokio::sync::mpsc::Sender<tokio::sync::oneshot::Sender<usize>>,
+    /// Shared with the pool's `TrackedConsumer` entry. Stream loop
+    /// flips it to `DISABLED` on merge entry so the pool stops
+    /// targeting this consumer.
+    reclaimer_state: Arc<std::sync::atomic::AtomicU8>,
 }
 
 #[async_trait::async_trait]
@@ -94,6 +98,10 @@ impl MemoryReclaimer for ExternalSorterReclaimer {
         }
         Ok(resp_rx.await.unwrap_or(0))
     }
+
+    fn reclaimer_state(&self) -> Option<Arc<std::sync::atomic::AtomicU8>> {
+        Some(Arc::clone(&self.reclaimer_state))
+    }
 }
 
 struct ExternalSorterMetrics {
@@ -1282,8 +1290,14 @@ impl ExecutionPlan for SortExec {
                 // Spill-request channel; drained by the stream loop below.
                 let (reclaim_tx, mut reclaim_rx) =
                     tokio::sync::mpsc::channel::<tokio::sync::oneshot::Sender<usize>>(4);
+                let state = Arc::new(std::sync::atomic::AtomicU8::new(
+                    reclaimer_state::AVAILABLE,
+                ));
                 let reclaimer: Arc<dyn MemoryReclaimer> =
-                    Arc::new(ExternalSorterReclaimer { tx: reclaim_tx });
+                    Arc::new(ExternalSorterReclaimer {
+                        tx: reclaim_tx,
+                        reclaimer_state: Arc::clone(&state),
+                    });
 
                 let mut sorter = ExternalSorter::new(
                     partition,
@@ -1338,7 +1352,17 @@ impl ExecutionPlan for SortExec {
                                 }
                             }
                         }
-                        // Late reclaim requests now resolve to Ok(0).
+                        // Sticky-disable so concurrent `try_grow_async`
+                        // callers stop targeting this consumer once we
+                        // enter the merge phase. Set before dropping
+                        // the receiver to close any window where the
+                        // pool would observe `AVAILABLE` after the
+                        // channel is gone (and hence get `Ok(0)` from a
+                        // wasted `reclaim`).
+                        state.store(
+                            reclaimer_state::DISABLED,
+                            std::sync::atomic::Ordering::Release,
+                        );
                         drop(reclaim_rx);
                         sorter.sort().await
                     })