Struct QueryMemoryTracker 

pub struct QueryMemoryTracker {
    pub(crate) current: Arc<AtomicUsize>,
    pub(crate) limit: usize,
    pub(crate) standard_tier_memory_fraction: f64,
    pub(crate) privileged_count: Arc<AtomicUsize>,
    pub(crate) privileged_slots: usize,
    pub(crate) on_update: Option<Arc<dyn Fn(usize) + Send + Sync>>,
    pub(crate) on_reject: Option<Arc<dyn Fn() + Send + Sync>>,
}

Memory tracker for RecordBatch streams. Clone to share the same limit across queries.

Implements a two-tier memory allocation strategy:

• Privileged tier: First N streams (default: 20) can use up to the full memory limit
• Standard tier: Remaining streams are restricted to a fraction of the limit (default: 70%)
• Privilege is granted on a first-come-first-served basis
• The configured limit is an absolute hard cap - no stream can exceed it

Fields

current: Arc<AtomicUsize>

limit: usize

standard_tier_memory_fraction: f64

privileged_count: Arc<AtomicUsize>

privileged_slots: usize

on_update: Option<Arc<dyn Fn(usize) + Send + Sync>>

on_reject: Option<Arc<dyn Fn() + Send + Sync>>

Implementations

impl QueryMemoryTracker

pub(crate) const DEFAULT_PRIVILEGED_SLOTS: usize = 20usize

pub(crate) const DEFAULT_PRIVILEGED_TIER_RATIO: f64 = 0.69999999999999996f64

pub fn new(limit: usize, max_concurrent_queries: usize) -> Self

Create a new memory tracker with the given limit and max_concurrent_queries. Calculates privileged slots as 70% of max_concurrent_queries (or 20 if max_concurrent_queries is 0).

Arguments

• limit - Maximum memory usage in bytes (hard limit for all streams). 0 means unlimited.
• max_concurrent_queries - Maximum number of concurrent queries (0 = unlimited).

pub fn with_privileged_slots(limit: usize, privileged_slots: usize) -> Self

Create a new memory tracker with custom privileged slots limit.

pub fn with_config( limit: usize, privileged_slots: usize, standard_tier_memory_fraction: f64, ) -> Self

Create a new memory tracker with full configuration.

Arguments

• limit - Maximum memory usage in bytes (hard limit for all streams). 0 means unlimited.
• privileged_slots - Maximum number of streams that can get privileged status.
• standard_tier_memory_fraction - Memory fraction for standard-tier streams (range: [0.0, 1.0]).

Panics

Panics if standard_tier_memory_fraction is not in the range [0.0, 1.0].

pub fn register_permit(&self) -> MemoryPermit

Register a new permit for memory tracking. The first privileged_slots permits get privileged status automatically. The returned permit can be shared across multiple streams of the same query.

pub fn with_on_update<F>(self, on_update: F) -> Self

where F: Fn(usize) + Send + Sync + 'static,

Set a callback to be called whenever the usage changes successfully. The callback receives the new total usage in bytes.

Note

The callback is called after both successful track() and release() operations. It is called even when limit == 0 (unlimited mode) to track actual usage.

pub fn with_on_reject<F>(self, on_reject: F) -> Self

where F: Fn() + Send + Sync + 'static,

Set a callback to be called when memory allocation is rejected.

Note

This is only called when track() fails due to exceeding the limit. It is never called when limit == 0 (unlimited mode).

pub fn current(&self) -> usize

Get the current memory usage in bytes.

pub(crate) fn calculate_privileged_slots(max_concurrent_queries: usize) -> usize

pub(crate) fn track_internal( &self, additional: usize, is_privileged: bool, stream_tracked: usize, ) -> Result<()>

Internal method to track additional memory usage.

Called by MemoryPermit::track(). Use MemoryPermit::track() instead of calling this directly.

pub fn release(&self, amount: usize)

Release tracked memory.

Arguments

• amount - Amount of memory to release in bytes

Trait Implementations

impl Clone for QueryMemoryTracker

fn clone(&self) -> QueryMemoryTracker

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for QueryMemoryTracker

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.