port reclaimer_state tri-state from old reclaimer-sortmerge-hashagg branch

The previous design relied solely on the close()+drain pattern to tell
the pool "this reclaimer can no longer free memory" — but the pool only
discovers that by invoking reclaim() and getting Ok(0) back. Snapshot,
sort by size, in-flight CAS, await reclaim, observe 0 — all wasted
per dead victim. Worse: dead reclaimers with high `reserved` sort to
the top, so they're picked first; live ones at lower reserved get
deprioritized. The dump's "reclaimable_sum" was inflated by these
phantom entries.

Port the tri-state mechanism from `reclaimer-sortmerge-hashagg`:

- `reclaimer::reclaimer_state` module: AVAILABLE / IN_FLIGHT / DISABLED.
- `MemoryReclaimer::reclaimer_state() -> Option<Arc<AtomicU8>>`: default
  None; implementations that participate in walks return Some(arc),
  shared with their owning operator.
- `TrackedConsumer.reclaim_in_flight: AtomicBool` removed; replaced with
  `TrackedConsumer.reclaimer_state: Arc<AtomicU8>` sourced from the
  reclaimer (or pool-allocated to AVAILABLE for in-flight dedup only).
- `ReclaimerStateGuard`: RAII CAS AVAILABLE -> IN_FLIGHT; Drop only
  restores AVAILABLE if state is still IN_FLIGHT (leaves DISABLED).
- `walk_reclaimers`: filter out DISABLED at snapshot time (logged as
  `skipped_disabled`); per-candidate CAS via ReclaimerStateGuard.
- `report_all_consumers`: three-way split — live-reclaimer (genuinely
  reclaimable), disabled-reclaimer (sticky-retired, dead bytes), and
  no-reclaimer (never reclaimable).

Operators sticky-set DISABLED right BEFORE closing reclaim_rx, so any
walk that snapshots after the store filters them at selection time:

- `ExternalSorterReclaimer` (sort.rs): in the post-input-loop block,
  store DISABLED then reclaim_rx.close() then drain.
- `HashAggReclaimer` (row_hash.rs): in the streaming-merge transition
  block where reclaim_rx is taken+closed.
- `MaterializingSmjReclaimer` (sort_merge_join): in the Exhausted state
  handler where reclaim_rx is taken+closed.

Q18-relevant effect: at OOM time, the consumer dump now correctly
separates the ~7 GiB of "appears reclaimable but actually returns 0"
bytes from the actually-reclaimable bytes, and walks no longer waste
round-trips on dead victims.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: JanKaul

datafusion/execution/src/memory_pool/mod.rs

@@ -40,7 +40,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

@@ -17,7 +17,7 @@
 
 use crate::memory_pool::{
     MemoryConsumer, MemoryLimit, MemoryPool, MemoryReclaimer, MemoryReservation,
-    human_readable_size,
+    human_readable_size, reclaimer_state,
 };
 use datafusion_common::HashMap;
 use datafusion_common::{DataFusionError, Result, resources_datafusion_err};
@@ -29,7 +29,7 @@ use std::pin::Pin;
 use std::sync::Arc;
 use std::{
     num::NonZeroUsize,
-    sync::atomic::{AtomicBool, AtomicUsize, Ordering},
+    sync::atomic::{AtomicU8, AtomicUsize, Ordering},
 };
 use tokio::sync::Semaphore;
 
@@ -332,13 +332,55 @@ struct TrackedConsumer {
     /// Hook the pool walks on `try_grow_async` failure. `None` when
     /// the consumer registered without a reclaimer.
     reclaimer: Option<Arc<dyn MemoryReclaimer>>,
-    /// Set while a reclaim walk is currently calling this consumer's
-    /// [`MemoryReclaimer::reclaim`]. Concurrent walks check-and-swap
-    /// before invoking the reclaimer so that at most one in-flight
-    /// reclaim per consumer happens at a time, even when the pool
-    /// allows multiple walks concurrently. Skipped candidates fall
-    /// through to the next victim in their list.
-    reclaim_in_flight: AtomicBool,
+    /// Shared tri-state flag controlling whether this reclaimer is
+    /// pickable by the walk. Values are defined in
+    /// [`reclaimer_state`]. The pool flips `AVAILABLE` ↔ `IN_FLIGHT`
+    /// for cross-walk dedup; the reclaimer's owner may sticky-set
+    /// `DISABLED` once it can no longer free memory (e.g. on entering
+    /// the non-reclaimable merge phase). If the reclaimer didn't
+    /// expose its own flag via [`MemoryReclaimer::reclaimer_state`],
+    /// the pool allocates a private `AVAILABLE`-initialised flag used
+    /// only for in-flight dedup.
+    reclaimer_state: Arc<AtomicU8>,
+}
+
+/// RAII guard for the [`IN_FLIGHT`] slot of a [`TrackedConsumer`]'s
+/// `reclaimer_state` flag. On `Drop` it only restores `AVAILABLE` if
+/// the state is still `IN_FLIGHT` — leaves a sticky `DISABLED` alone
+/// so a victim can retire mid-walk without being re-armed.
+///
+/// [`IN_FLIGHT`]: reclaimer_state::IN_FLIGHT
+struct ReclaimerStateGuard {
+    flag: Arc<AtomicU8>,
+}
+
+impl Drop for ReclaimerStateGuard {
+    fn drop(&mut self) {
+        let _ = self.flag.compare_exchange(
+            reclaimer_state::IN_FLIGHT,
+            reclaimer_state::AVAILABLE,
+            Ordering::AcqRel,
+            Ordering::Relaxed,
+        );
+    }
+}
+
+impl ReclaimerStateGuard {
+    /// Try to transition the flag from `AVAILABLE` to `IN_FLIGHT`.
+    /// Fails on contention (another walker won the CAS) 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),
+        })
+    }
 }
 
 impl TrackedConsumer {
@@ -587,7 +629,16 @@ impl<I: MemoryPool> MemoryPool for TrackConsumersPool<I> {
                 reserved: Default::default(),
                 peak: Default::default(),
                 reclaimer: consumer.reclaimer().cloned(),
-                reclaim_in_flight: AtomicBool::new(false),
+                // Source the tri-state flag from the reclaimer if it
+                // exposes one (so the operator can sticky-set
+                // DISABLED). Otherwise allocate a private AVAILABLE
+                // flag used only for in-flight dedup across walks.
+                reclaimer_state: consumer
+                    .reclaimer()
+                    .and_then(|r| r.reclaimer_state())
+                    .unwrap_or_else(|| {
+                        Arc::new(AtomicU8::new(reclaimer_state::AVAILABLE))
+                    }),
             },
         );
 
@@ -777,13 +828,22 @@ impl<I: MemoryPool> TrackConsumersPool<I> {
         min_reserved: usize,
     ) -> usize {
         // Snapshot candidates while holding the read lock; release
-        // before any await. We carry the `cid` and name along so we
-        // can flip the per-victim in-flight flag without re-locking
-        // and log meaningfully.
-        let (mut candidates, total_eligible) = {
+        // before any await. Carry the shared `reclaimer_state` flag
+        // along so we can do the CAS without re-locking, and skip
+        // sticky-DISABLED reclaimers at filter time instead of
+        // wasting a `reclaim()` call that would return `Ok(0)`.
+        let (mut candidates, total_eligible, skipped_disabled) = {
             let guard = self.tracked_consumers.read();
             let mut total = 0usize;
-            let cands: Vec<(usize, String, usize, Arc<dyn MemoryReclaimer>, i32)> = guard
+            let mut skipped_disabled = 0usize;
+            let cands: Vec<(
+                usize,
+                String,
+                usize,
+                Arc<dyn MemoryReclaimer>,
+                i32,
+                Arc<AtomicU8>,
+            )> = guard
                 .iter()
                 .filter_map(|(cid, tc)| {
                     if Some(*cid) == exclude_consumer_id {
@@ -794,17 +854,30 @@ impl<I: MemoryPool> TrackConsumersPool<I> {
                     if reserved <= min_reserved {
                         return None;
                     }
+                    // Skip sticky-DISABLED reclaimers (operator has
+                    // retired itself). IN_FLIGHT is also non-AVAILABLE
+                    // but we still record those at the per-candidate
+                    // CAS site below for the `skipped_in_flight`
+                    // counter — they may transition back to AVAILABLE
+                    // by the time we reach them in the iteration.
+                    if tc.reclaimer_state.load(Ordering::Acquire)
+                        == reclaimer_state::DISABLED
+                    {
+                        skipped_disabled += 1;
+                        return None;
+                    }
                     total = total.saturating_add(reserved);
                     Some((
                         *cid,
                         tc.name.clone(),
                         reserved,
                         Arc::clone(reclaimer),
                         reclaimer.priority(),
+                        Arc::clone(&tc.reclaimer_state),
                     ))
                 })
                 .collect();
-            (cands, total)
+            (cands, total, skipped_disabled)
         };
         // Order: priority DESC, then reserved-size DESC. No cap —
         // walk every eligible candidate; the in-flight flag handles
@@ -813,43 +886,33 @@ impl<I: MemoryPool> TrackConsumersPool<I> {
 
         debug!(
             "[reclaim-walk] excl={:?} min_reserved={} target={} candidates={} \
-             eligible_bytes={} pool_reserved={}",
+             eligible_bytes={} skipped_disabled={} pool_reserved={}",
             exclude_consumer_id,
             min_reserved,
             target,
             candidates.len(),
             total_eligible,
+            skipped_disabled,
             self.inner.reserved(),
         );
 
         let mut total_freed: usize = 0;
         let mut still_needed = target;
         let mut skipped_in_flight = 0usize;
-        for (cid, name, _size, reclaimer, _prio) in candidates {
+        for (cid, name, _size, reclaimer, _prio, flag) in candidates {
             if still_needed == 0 {
                 break;
             }
-            // Atomic check-and-set: if another walk is already
-            // reclaiming this victim, skip and try the next one. We
-            // re-fetch the entry by cid (held read lock briefly).
-            let claimed = {
-                let guard = self.tracked_consumers.read();
-                guard
-                    .get(&cid)
-                    .is_some_and(|tc| !tc.reclaim_in_flight.swap(true, Ordering::AcqRel))
-            };
-            if !claimed {
+            // CAS AVAILABLE → IN_FLIGHT. Fails if (a) another walker
+            // got there first (in-flight) or (b) the owner sticky-set
+            // DISABLED between snapshot and now. Either way, skip.
+            // Guard's Drop restores AVAILABLE on scope exit (leaving
+            // DISABLED alone).
+            let Some(_g) = ReclaimerStateGuard::try_acquire(&flag) else {
                 skipped_in_flight += 1;
                 continue;
-            }
+            };
             let result = reclaimer.reclaim(still_needed).await;
-            // Release the in-flight flag whether reclaim succeeded or
-            // not. If the consumer was deregistered mid-walk
-            // `get(&cid)` returns None; that's fine — the flag was on
-            // a now-dropped value.
-            if let Some(tc) = self.tracked_consumers.read().get(&cid) {
-                tc.reclaim_in_flight.store(false, Ordering::Release);
-            }
             match result {
                 Ok(freed) => {
                     debug!(
@@ -875,26 +938,33 @@ impl<I: MemoryPool> TrackConsumersPool<I> {
     }
 
     /// Diagnostic: format ALL tracked consumers with non-zero current
-    /// reservations, grouped by whether they have a reclaimer or not.
+    /// reservations, split by reclaimer state so dumps distinguish
+    /// truly-reclaimable bytes from phantom (sticky-DISABLED) ones.
     /// Useful at OOM time to see the full pool composition beyond
     /// what [`Self::report_top`] shows.
     pub fn report_all_consumers(&self) -> String {
         let guard = self.tracked_consumers.read();
-        let mut with_reclaimer: Vec<(usize, String, usize, usize)> = Vec::new();
+        let mut live: Vec<(usize, String, usize, usize)> = Vec::new();
+        let mut disabled: Vec<(usize, String, usize, usize)> = Vec::new();
         let mut without_reclaimer: Vec<(usize, String, usize, usize)> = Vec::new();
         for (cid, tc) in guard.iter() {
             let reserved = tc.reserved();
             if reserved == 0 {
                 continue;
             }
             let entry = (*cid, tc.name.clone(), reserved, tc.peak());
-            if tc.reclaimer.is_some() {
-                with_reclaimer.push(entry);
-            } else {
+            if tc.reclaimer.is_none() {
                 without_reclaimer.push(entry);
+            } else if tc.reclaimer_state.load(Ordering::Acquire)
+                == reclaimer_state::DISABLED
+            {
+                disabled.push(entry);
+            } else {
+                live.push(entry);
             }
         }
-        with_reclaimer.sort_by_key(|e| std::cmp::Reverse(e.2));
+        live.sort_by_key(|e| std::cmp::Reverse(e.2));
+        disabled.sort_by_key(|e| std::cmp::Reverse(e.2));
         without_reclaimer.sort_by_key(|e| std::cmp::Reverse(e.2));
         let fmt = |v: &[(usize, String, usize, usize)]| -> (String, usize) {
             let sum = v.iter().map(|e| e.2).sum::<usize>();
@@ -911,15 +981,22 @@ impl<I: MemoryPool> TrackConsumersPool<I> {
                 .join("\n");
             (body, sum)
         };
-        let (with_body, with_sum) = fmt(&with_reclaimer);
+        let (live_body, live_sum) = fmt(&live);
+        let (disabled_body, disabled_sum) = fmt(&disabled);
         let (without_body, without_sum) = fmt(&without_reclaimer);
         format!(
-            "[consumer-dump] pool_reserved={} reclaimable_sum={} ({} consumers):\n{}\n\
+            "[consumer-dump] pool_reserved={}\n\
+             reclaimable_sum={} ({} consumers):\n{}\n\
+             disabled_reclaimer_sum={} ({} consumers — sticky-retired, \
+             pool can no longer free these):\n{}\n\
              non-reclaimable_sum={} ({} consumers):\n{}",
             human_readable_size(self.inner.reserved()),
-            human_readable_size(with_sum),
-            with_reclaimer.len(),
-            with_body,
+            human_readable_size(live_sum),
+            live.len(),
+            live_body,
+            human_readable_size(disabled_sum),
+            disabled.len(),
+            disabled_body,
             human_readable_size(without_sum),
             without_reclaimer.len(),
             without_body,

datafusion/execution/src/memory_pool/reclaimer.rs

@@ -22,6 +22,29 @@
 
 use datafusion_common::Result;
 use std::fmt::Debug;
+use std::sync::Arc;
+use std::sync::atomic::AtomicU8;
+
+/// Encoded values stored in the shared [`reclaimer_state`] tri-state
+/// flag. The flag is `Arc<AtomicU8>` so the operator (owner) and the
+/// pool (selector) can observe transitions atomically.
+///
+/// [`reclaimer_state`]: MemoryReclaimer::reclaimer_state
+pub mod reclaimer_state {
+    /// Reclaimer is idle and may be selected as a walk victim.
+    pub const AVAILABLE: u8 = 0;
+    /// A pool walk is currently driving `reclaim` on this reclaimer.
+    /// Other concurrent walks should skip it to avoid serializing on
+    /// the same victim's mpsc/select! arm.
+    pub const IN_FLIGHT: u8 = 1;
+    /// Reclaimer has been sticky-retired by its owner — typically when
+    /// the operator transitions past the phase where it can free
+    /// memory (e.g., a sort entering its non-reclaimable merge phase).
+    /// Once set, never returns to `AVAILABLE`; walks observe it
+    /// immediately and skip without invoking `reclaim` (which would
+    /// just return `Ok(0)` and waste a round-trip).
+    pub const DISABLED: u8 = 2;
+}
 
 /// Hook attached to a [`MemoryConsumer`] via
 /// [`MemoryConsumer::with_reclaimer`]. On
@@ -64,4 +87,21 @@ pub trait MemoryReclaimer: Send + Sync + Debug {
     fn priority(&self) -> i32 {
         0
     }
+
+    /// Optional shared flag the pool consults before invoking
+    /// [`Self::reclaim`]. Implementations that participate in walks
+    /// return `Some(arc)` and share the same `Arc<AtomicU8>` with the
+    /// owning operator, so the operator can sticky-set
+    /// [`reclaimer_state::DISABLED`] when it can no longer free memory
+    /// (e.g., after closing its reclaim channel on the way into the
+    /// merge phase). The pool then immediately stops picking this
+    /// reclaimer as a victim instead of wasting a `reclaim()` call
+    /// that would return `Ok(0)`.
+    ///
+    /// Implementations that don't expose this flag (return `None`)
+    /// get a pool-allocated `AVAILABLE` flag used only for in-flight
+    /// dedup across concurrent walks.
+    fn reclaimer_state(&self) -> Option<Arc<AtomicU8>> {
+        None
+    }
 }