use reclaimer handle

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, reclaimer_state};
+pub use reclaimer::{MemoryReclaimer, ReclaimerHandle};
 
 /// Tracks and potentially limits memory use across operators during execution.
 ///
@@ -264,9 +264,11 @@ pub struct MemoryConsumer {
     name: String,
     can_spill: bool,
     id: usize,
-    /// Reclaimer collected by reclaim-aware pools at register time. Not
-    /// part of consumer identity (excluded from `Eq`/`Hash`).
-    reclaimer: Option<Arc<dyn MemoryReclaimer>>,
+    /// Reclaim hook collected by reclaim-aware pools at register time.
+    /// The handle is the operator's control over eligibility (sticky
+    /// disable); the trait object is the action. Not part of consumer
+    /// identity (excluded from `Eq`/`Hash`).
+    reclaim: Option<(Arc<dyn MemoryReclaimer>, ReclaimerHandle)>,
 }
 
 impl PartialEq for MemoryConsumer {
@@ -305,37 +307,50 @@ impl MemoryConsumer {
             name: name.into(),
             can_spill: false,
             id: Self::new_unique_id(),
-            reclaimer: None,
+            reclaim: None,
         }
     }
 
     /// Clone this [`MemoryConsumer`] with a new unique id.
     ///
-    /// Drops any attached reclaimer: it is bound to the original operator's
-    /// state and would target the wrong owner under a new id (and bypass
-    /// the id-keyed requestor-self-skip in `try_grow_async`).
+    /// Drops any attached reclaimer/handle: they are bound to the
+    /// original operator's state and would target the wrong owner
+    /// under a new id (and bypass the id-keyed requestor-self-skip in
+    /// `try_grow_async`).
     pub fn clone_with_new_id(&self) -> Self {
         Self {
             name: self.name.clone(),
             can_spill: self.can_spill,
             id: Self::new_unique_id(),
-            reclaimer: None,
+            reclaim: None,
         }
     }
 
-    /// Attach a [`MemoryReclaimer`] and mark this consumer spill-capable.
-    /// Pools without reclaim support ignore the reclaimer.
-    pub fn with_reclaimer(self, reclaimer: Arc<dyn MemoryReclaimer>) -> Self {
+    /// Attach a [`MemoryReclaimer`] with its control [`ReclaimerHandle`]
+    /// and mark this consumer spill-capable. Pools without reclaim
+    /// support ignore both. The operator keeps a clone of `handle` to
+    /// later [`ReclaimerHandle::disable`] itself once it can no longer
+    /// free memory.
+    pub fn with_reclaimer(
+        self,
+        reclaimer: Arc<dyn MemoryReclaimer>,
+        handle: ReclaimerHandle,
+    ) -> Self {
         Self {
             can_spill: true,
-            reclaimer: Some(reclaimer),
+            reclaim: Some((reclaimer, handle)),
             ..self
         }
     }
 
     /// Returns the attached [`MemoryReclaimer`], if any.
     pub fn reclaimer(&self) -> Option<&Arc<dyn MemoryReclaimer>> {
-        self.reclaimer.as_ref()
+        self.reclaim.as_ref().map(|(r, _)| r)
+    }
+
+    /// Returns the attached [`ReclaimerHandle`], if any.
+    pub fn reclaimer_handle(&self) -> Option<&ReclaimerHandle> {
+        self.reclaim.as_ref().map(|(_, h)| h)
     }
 
     /// Return the unique id of this [`MemoryConsumer`]

datafusion/execution/src/memory_pool/pool.rs

@@ -18,7 +18,7 @@
 use crate::memory_pool::reclaimer::reclaimer_state;
 use crate::memory_pool::{
     MemoryConsumer, MemoryLimit, MemoryPool, MemoryReclaimer, MemoryReservation,
-    human_readable_size,
+    ReclaimerHandle, human_readable_size,
 };
 use datafusion_common::HashMap;
 use datafusion_common::{DataFusionError, Result, resources_datafusion_err};
@@ -343,13 +343,13 @@ struct TrackedConsumer {
     reserved: AtomicUsize,
     peak: AtomicUsize,
     reclaimer: Option<Arc<dyn MemoryReclaimer>>,
-    /// 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>,
+    /// Eligibility handle for [`reclaimer`]. The pool flips
+    /// `AVAILABLE` ↔ `IN_FLIGHT` on its internal cell for dedup; the
+    /// reclaimer's owner sticky-sets `DISABLED` via [`ReclaimerHandle`].
+    /// For consumers without a reclaimer the pool still allocates a
+    /// fresh handle so the `self_guard` path in `try_grow_async` is
+    /// uniform (cheap: one `Arc<AtomicU8>`).
+    handle: ReclaimerHandle,
 }
 
 /// RAII guard for the [`IN_FLIGHT`] slot of a [`TrackedConsumer`]'s
@@ -628,14 +628,15 @@ impl<I: MemoryPool> MemoryPool for TrackConsumersPool<I> {
         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)));
+        // Use the operator-supplied handle when one is attached, so a
+        // sticky `DISABLED` flip on the operator side is visible to
+        // the pool on the next filter pass. For consumers without a
+        // reclaimer, allocate a private handle so the in-flight guard
+        // path stays uniform.
+        let handle = consumer
+            .reclaimer_handle()
+            .cloned()
+            .unwrap_or_else(ReclaimerHandle::new);
 
         let mut guard = self.tracked_consumers.write();
         let existing = guard.insert(
@@ -646,7 +647,7 @@ impl<I: MemoryPool> MemoryPool for TrackConsumersPool<I> {
                 reserved: Default::default(),
                 peak: Default::default(),
                 reclaimer,
-                reclaimer_state: state,
+                handle,
             },
         );
 
@@ -738,23 +739,36 @@ impl<I: MemoryPool> MemoryPool for TrackConsumersPool<I> {
                 .tracked_consumers
                 .read()
                 .get(&requestor_id)
-                .and_then(|tc| ReclaimerStateGuard::try_acquire(&tc.reclaimer_state));
+                .and_then(|tc| ReclaimerStateGuard::try_acquire(tc.handle.state()));
 
             let mut retries: usize = 0;
             loop {
-                // Snapshot reclaimers. Only consumers strictly larger than
-                // the requestor are eligible: smaller-or-equal siblings would
-                // free less than the requestor itself can, so the requestor
-                // should self-spill instead. This rule also breaks the
-                // 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. Filter out anyone
-                // whose `reclaimer_state` flag is not `AVAILABLE` (in-flight or
-                // sticky-disabled). Also count IN_FLIGHT siblings so we know
-                // whether to wait briefly for them to finish before giving up.
-                // Drop the read guard before awaiting any reclaim.
-                let requestor_reserved = {
+                // Snapshot reclaimers. Eligibility ranks by *reclaimable*
+                // bytes (what the reclaimer believes it can free now),
+                // not by total consumer reservation: a sorter whose
+                // buffer just spilled may still carry a non-zero merge
+                // reservation, but its reclaimable size is ~0 and
+                // picking it would burn a round trip for nothing.
+                // Reclaimers that don't report a precise bound fall
+                // back to their tracked `reserved` as a conservative
+                // upper bound.
+                //
+                // Only consumers with *strictly more* reclaimable bytes
+                // than the requestor are eligible: smaller-or-equal
+                // siblings would free less than the requestor itself
+                // could, so the requestor should self-spill instead.
+                // This rule also breaks the mutual-reclaim cycle (A
+                // targets B while B targets A) — at most one side of
+                // any pair can hold strictly more, so the other side
+                // has no candidates and surfaces an error for the
+                // caller's self-spill fallback.
+                //
+                // Filter out anyone whose handle is not `AVAILABLE`
+                // (in-flight or sticky-disabled). Also count IN_FLIGHT
+                // siblings so we know whether to wait briefly for them
+                // to finish before giving up. Drop the read guard
+                // before awaiting any reclaim.
+                let requestor_reclaimable = {
                     let guard = self.tracked_consumers.read();
                     guard
                         .get(&requestor_id)
@@ -776,26 +790,27 @@ impl<I: MemoryPool> MemoryPool for TrackConsumersPool<I> {
                             }
                             // Track in-flight siblings (any size) so we can
                             // decide whether a retry has any chance of helping.
-                            let state = tc.reclaimer_state.load(Ordering::Acquire);
+                            let state = tc.handle.state().load(Ordering::Acquire);
                             if state == reclaimer_state::IN_FLIGHT {
                                 in_flight_seen += 1;
                             }
                             let reclaimer = tc.reclaimer.as_ref()?;
-                            if tc.reserved() <= requestor_reserved {
+                            let reclaimable = tc.reserved();
+                            if reclaimable <= requestor_reclaimable {
                                 return None;
                             }
                             if state != reclaimer_state::AVAILABLE {
                                 return None;
                             }
                             Some((
-                                tc.reserved(),
+                                reclaimable,
                                 Arc::clone(reclaimer),
-                                Arc::clone(&tc.reclaimer_state),
+                                Arc::clone(tc.handle.state()),
                             ))
                         })
                         .collect()
                 };
-                // Order: priority desc, then reservation size desc.
+                // Order: priority desc, then reclaimable bytes desc.
                 candidates.sort_by(|(lr, l, _), (rr, r, _)| {
                     r.priority().cmp(&l.priority()).then_with(|| rr.cmp(lr))
                 });

datafusion/execution/src/memory_pool/reclaimer.rs

@@ -23,12 +23,10 @@
 use datafusion_common::Result;
 use std::fmt::Debug;
 use std::sync::Arc;
-use std::sync::atomic::AtomicU8;
+use std::sync::atomic::{AtomicU8, Ordering};
 
-/// Encoded values stored in the [`reclaimer_state`] tri-state.
-///
-/// [`reclaimer_state`]: MemoryReclaimer::reclaimer_state
-pub mod reclaimer_state {
+/// Encoded values stored in a [`ReclaimerHandle`]'s tri-state cell.
+pub(crate) 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.
@@ -39,6 +37,55 @@ pub mod reclaimer_state {
     pub const DISABLED: u8 = 2;
 }
 
+/// Operator-side control over a reclaimer's eligibility in the pool.
+///
+/// Created by the operator and passed alongside a [`MemoryReclaimer`]
+/// to [`MemoryConsumer::with_reclaimer`]. The pool stores the same
+/// handle on its tracking entry, so any [`Self::disable`] call is
+/// observed on the next reclaim walk.
+///
+/// The internal `IN_FLIGHT` slot is managed by the pool for dedup and
+/// is opaque to operators; operators only flip `AVAILABLE` → `DISABLED`.
+///
+/// [`MemoryConsumer::with_reclaimer`]: super::MemoryConsumer::with_reclaimer
+#[derive(Debug, Clone)]
+pub struct ReclaimerHandle {
+    state: Arc<AtomicU8>,
+}
+
+impl ReclaimerHandle {
+    /// Create a new handle in the `AVAILABLE` state.
+    pub fn new() -> Self {
+        Self {
+            state: Arc::new(AtomicU8::new(reclaimer_state::AVAILABLE)),
+        }
+    }
+
+    /// Sticky-disable this reclaimer. After this call the pool will
+    /// stop selecting the associated consumer as a reclaim victim.
+    /// Idempotent.
+    pub fn disable(&self) {
+        self.state
+            .store(reclaimer_state::DISABLED, Ordering::Release);
+    }
+
+    /// Returns `true` once [`Self::disable`] has been called.
+    pub fn is_disabled(&self) -> bool {
+        self.state.load(Ordering::Acquire) == reclaimer_state::DISABLED
+    }
+
+    /// Pool-internal access to the shared state cell.
+    pub(crate) fn state(&self) -> &Arc<AtomicU8> {
+        &self.state
+    }
+}
+
+impl Default for ReclaimerHandle {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 /// Hook attached to a [`MemoryConsumer`] via
 /// [`MemoryConsumer::with_reclaimer`]. On
 /// [`MemoryPool::try_grow_async`] failure the pool walks registered
@@ -54,18 +101,18 @@ pub mod reclaimer_state {
 /// - Not capture `Arc<MemoryReservation>` / `Arc<MemoryConsumer>`
 ///   (creates a cycle that blocks `unregister`); use channels or `Weak`.
 ///
+/// Eligibility (whether the pool may pick this reclaimer at all) is
+/// controlled out-of-band via the [`ReclaimerHandle`] passed at
+/// registration. The trait itself is narrowly about "how many bytes
+/// can you free, and please free some now."
+///
 /// [`MemoryConsumer`]: super::MemoryConsumer
 /// [`MemoryConsumer::with_reclaimer`]: super::MemoryConsumer::with_reclaimer
 /// [`MemoryPool::try_grow_async`]: super::MemoryPool::try_grow_async
 /// [`MemoryReservation::shrink`]: super::MemoryReservation::shrink
 /// [`MemoryReservation::free`]: super::MemoryReservation::free
 #[async_trait::async_trait]
 pub trait MemoryReclaimer: Send + Sync + Debug {
-    /// Upper bound on bytes this reclaimer can free. `None` = unknown.
-    fn reclaimable_bytes(&self) -> Option<usize> {
-        None
-    }
-
     /// Free up to `target` bytes; return the amount actually released.
     /// See trait-level contract.
     async fn reclaim(&self, target: usize) -> Result<usize>;
@@ -74,15 +121,4 @@ 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
-    }
 }