micromegas_analytics/lakehouse/
scoped_memory_pool.rs

1//! Per-query [`datafusion::execution::memory_pool::MemoryPool`] wrapper that adds accounting on top of a shared pool.
2//!
3//! `ScopedMemoryPool` delegates every operation to the wrapped `inner` pool, so the
4//! budget, spilling behavior, and top-consumer reporting (when `inner` is a
5//! `TrackConsumersPool`) stay exactly as they were. The only thing it adds is its own
6//! `current`/`peak` counters, which track just the reservations that were grown through
7//! this particular instance. Because a `MemoryReservation` is permanently bound to the
8//! `Arc<dyn MemoryPool>` it registered with (see `MemoryConsumer::register`), wrapping the
9//! shared pool once per query and handing that wrapper to the query's `RuntimeEnv` is
10//! enough to isolate one query's peak from another's, with no ambient context or naming
11//! scheme involved.
12use datafusion::error::Result;
13use datafusion::execution::memory_pool::{
14    MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation,
15};
16use std::sync::Arc;
17use std::sync::atomic::{AtomicUsize, Ordering};
18
19/// A [`MemoryPool`] that forwards every call to `inner`, while separately tracking the
20/// current and peak (high-water mark) reservation grown through this instance alone.
21#[derive(Debug)]
22pub struct ScopedMemoryPool {
23    inner: Arc<dyn MemoryPool>,
24    current: AtomicUsize,
25    peak: AtomicUsize,
26}
27
28impl ScopedMemoryPool {
29    /// Wraps `inner`, starting both counters at zero.
30    pub fn new(inner: Arc<dyn MemoryPool>) -> Self {
31        Self {
32            inner,
33            current: AtomicUsize::new(0),
34            peak: AtomicUsize::new(0),
35        }
36    }
37
38    /// Monotonic high-water mark of tracked reservation for this query, in bytes.
39    pub fn peak(&self) -> usize {
40        self.peak.load(Ordering::Relaxed)
41    }
42
43    /// Currently outstanding tracked reservation, in bytes; `0` at quiescence.
44    pub fn current(&self) -> usize {
45        self.current.load(Ordering::Relaxed)
46    }
47
48    /// Records a successful grow of `additional` bytes: bumps `current` and, if that
49    /// pushed it to a new high, `peak` along with it.
50    fn record_grow(&self, additional: usize) {
51        let cur = self.current.fetch_add(additional, Ordering::Relaxed) + additional;
52        self.peak.fetch_max(cur, Ordering::Relaxed);
53    }
54}
55
56impl MemoryPool for ScopedMemoryPool {
57    fn name(&self) -> &str {
58        "scoped"
59    }
60
61    // MUST forward: the trait's default `register`/`unregister` are no-ops, and
62    // `TrackConsumersPool`'s whole tracked-consumers map (and therefore the top-consumers
63    // text in its OOM error) is driven by them.
64    fn register(&self, consumer: &MemoryConsumer) {
65        self.inner.register(consumer);
66    }
67
68    fn unregister(&self, consumer: &MemoryConsumer) {
69        self.inner.unregister(consumer);
70    }
71
72    fn grow(&self, reservation: &MemoryReservation, additional: usize) {
73        self.inner.grow(reservation, additional);
74        self.record_grow(additional);
75    }
76
77    fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> {
78        self.inner.try_grow(reservation, additional)?; // `?` keeps counters untouched on failure
79        self.record_grow(additional);
80        Ok(())
81    }
82
83    fn shrink(&self, reservation: &MemoryReservation, shrink: usize) {
84        self.inner.shrink(reservation, shrink);
85        let prev = self.current.fetch_sub(shrink, Ordering::Relaxed);
86        debug_assert!(prev >= shrink, "scoped pool shrink underflow");
87    }
88
89    fn reserved(&self) -> usize {
90        self.inner.reserved() // global truth, unchanged semantics
91    }
92
93    fn memory_limit(&self) -> MemoryLimit {
94        self.inner.memory_limit()
95    }
96}
97
98impl std::fmt::Display for ScopedMemoryPool {
99    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
100        write!(
101            f,
102            "scoped(inner_pool: {}, peak: {})",
103            self.inner,
104            self.peak()
105        )
106    }
107}
micromegas_analytics/lakehouse/scoped_memory_pool.rs

micromegas_analytics/lakehouse/
scoped_memory_pool.rs
