Use resident_bytes instead of allocated_bytes in memory guard

The memory guard's should_override() used jemalloc's allocated_bytes
(live objects) to determine whether pool rejections were false positives.
Under concurrent load, jemalloc retains freed pages in thread caches
and arenas (dirty/muzzy decay), causing allocated to lag 10-20GB behind
resident (physical RSS). This gap allowed overrides to fire when the
system was near OOM.

Changes:
- should_override() now uses resident_bytes (physical RSS) for both
  admission and operator contexts
- New is_memory_pressured() function for proactive admission check
- query_budget: proactive RSS check at admission — when pool reserved
  >= admission threshold AND RSS confirms pressure, immediately reduce
  to min partitions (prevents concurrent burst)
- Jemalloc RSS is only consulted when pool accounting shows >= threshold
  utilization (no overhead on the happy path)

Behavior:
- 0-70% pool utilized: no jemalloc call, full partitions (unchanged)
- 70%+ pool AND RSS < threshold: override fires (stale accounting)
- 70%+ pool AND RSS >= threshold: reduce partitions / trigger spill

Signed-off-by: Bukhtawar Khan <bukhtawa@amazon.com>

Author: Bukhtawar

sandbox/plugins/analytics-backend-datafusion/rust/src/memory_guard.rs

@@ -83,6 +83,12 @@ pub fn get_thresholds() -> MemoryThresholds {
 /// Check whether jemalloc says physical memory has headroom, meaning the
 /// pool's rejection is a false positive (stale accounting).
 ///
+/// Uses `resident_bytes` (physical RSS) instead of `allocated_bytes` (live objects).
+/// `allocated_bytes` undercounts true memory pressure because jemalloc retains
+/// freed pages in thread caches and arenas (dirty/muzzy decay). Under concurrent
+/// workloads, the gap between allocated and resident can be 10-20GB, causing the
+/// override to fire when the system is actually near OOM.
+///
 /// Returns `true` if the override should fire (proceed despite pool rejection).
 /// Returns `false` if pressure is real or stats are unavailable.
 ///
@@ -95,8 +101,8 @@ pub fn should_override(pool_limit_bytes: usize, context: OverrideContext) -> boo
         return false;
     }
 
-    let allocated = native_bridge_common::allocator::allocated_bytes();
-    if allocated <= 0 {
+    let resident = native_bridge_common::allocator::resident_bytes();
+    if resident <= 0 {
         return false;
     }
 
@@ -106,7 +112,31 @@ pub fn should_override(pool_limit_bytes: usize, context: OverrideContext) -> boo
     };
 
     let threshold_bytes = (pool_limit_bytes as u64 * threshold_x1000 / 1000) as i64;
-    allocated < threshold_bytes
+    resident < threshold_bytes
+}
+
+/// Proactive admission check: returns `true` if jemalloc resident memory
+/// already exceeds the admission threshold (70% of pool limit by default).
+///
+/// Called BEFORE query execution (at budget acquisition) to reject or reduce
+/// concurrency early — before any hash table allocation occurs. This prevents
+/// the "20 queries all pass admission simultaneously" burst that causes OOM.
+///
+/// Cost: one `epoch.advance` + stat read (~1-5µs). Called once per query at
+/// admission, not per-batch.
+pub fn is_memory_pressured(pool_limit_bytes: usize) -> bool {
+    if pool_limit_bytes < MIN_POOL_FOR_OVERRIDE {
+        return false;
+    }
+
+    let resident = native_bridge_common::allocator::resident_bytes();
+    if resident <= 0 {
+        return false;
+    }
+
+    let threshold_x1000 = ADMISSION_THRESHOLD_X1000.load(Ordering::Acquire);
+    let threshold_bytes = (pool_limit_bytes as u64 * threshold_x1000 / 1000) as i64;
+    resident >= threshold_bytes
 }
 
 // ---------------------------------------------------------------------------
@@ -203,4 +233,76 @@ mod tests {
         assert!(!should_override(1_000_000, OverrideContext::Admission));
         assert!(!should_override(1_000_000, OverrideContext::Operator));
     }
+
+    #[test]
+    fn should_override_uses_resident_not_allocated() {
+        // With a large pool (1TB), resident will always be below threshold
+        // so override should fire (resident < threshold = "headroom available")
+        let large_pool = 1024 * 1024 * 1024 * 1024; // 1TB
+        // This test validates the function runs without error and uses resident.
+        // On a test process with < 1TB RSS, override should fire (we have headroom).
+        let result = should_override(large_pool, OverrideContext::Operator);
+        assert!(result, "With 1TB pool limit, resident should be well below threshold — override should fire");
+    }
+
+    #[test]
+    fn is_memory_pressured_false_for_large_pool() {
+        // With a 1TB pool, current process RSS is far below 70% → not pressured
+        let large_pool = 1024 * 1024 * 1024 * 1024; // 1TB
+        assert!(!is_memory_pressured(large_pool));
+    }
+
+    #[test]
+    fn is_memory_pressured_true_when_rss_exceeds_limit() {
+        // Set pool limit to something well below current process RSS.
+        // A Rust test process typically uses 50-200MB RSS, so a 20MB limit
+        // should always be exceeded.
+        let small_pool = 20 * 1024 * 1024; // 20MB — above MIN_POOL_FOR_OVERRIDE
+        let resident = native_bridge_common::allocator::resident_bytes();
+        if resident <= 0 {
+            return; // jemalloc not available
+        }
+        // Only assert if RSS is actually above 70% of 20MB = 14MB (which it will be)
+        if resident as usize > small_pool * 70 / 100 {
+            assert!(is_memory_pressured(small_pool));
+        }
+    }
+
+    #[test]
+    fn is_memory_pressured_skips_small_pools() {
+        assert!(!is_memory_pressured(1_000_000)); // 1MB — below MIN_POOL_FOR_OVERRIDE
+    }
+
+    #[test]
+    fn override_respects_operator_vs_admission_threshold() {
+        // Operator threshold (85%) is more permissive than admission (70%).
+        // For a pool where RSS is between 70% and 85%:
+        // - Admission override should NOT fire (RSS >= 70% threshold)
+        // - Operator override SHOULD fire (RSS < 85% threshold)
+        //
+        // We can't precisely control RSS in a unit test, but we can verify
+        // that the thresholds are read correctly by setting them and checking
+        // behavior with known pool sizes.
+        let resident = native_bridge_common::allocator::resident_bytes();
+        if resident <= 0 {
+            return; // jemalloc not available in this test env
+        }
+        let resident = resident as usize;
+
+        // Set pool limit so that resident is exactly between 70% and 85%
+        // pool = resident / 0.77 (midpoint) → resident/pool ≈ 77%
+        let pool_at_midpoint = (resident as f64 / 0.77) as usize;
+        if pool_at_midpoint < MIN_POOL_FOR_OVERRIDE {
+            return;
+        }
+
+        // At 77% utilization: admission (70%) should NOT override, operator (85%) SHOULD override
+        let admission_result = should_override(pool_at_midpoint, OverrideContext::Admission);
+        let operator_result = should_override(pool_at_midpoint, OverrideContext::Operator);
+
+        // admission: resident (77%) >= threshold (70%) → NOT below → override = false
+        assert!(!admission_result, "At 77% RSS, admission override should NOT fire (threshold 70%)");
+        // operator: resident (77%) < threshold (85%) → below → override = true
+        assert!(operator_result, "At 77% RSS, operator override SHOULD fire (threshold 85%)");
+    }
 }

sandbox/plugins/analytics-backend-datafusion/rust/src/query_budget.rs

@@ -191,6 +191,13 @@ pub fn acquire_budget_with_projection(
 ///
 /// Tries to reserve a phantom at the configured parallelism. If the pool
 /// rejects it, iteratively halves target_partitions until it fits.
+///
+/// **Proactive RSS check**: before attempting pool reservation, consults
+/// jemalloc's resident bytes. If physical memory already exceeds the admission
+/// threshold (default 70% of pool limit), immediately reduces partitions to
+/// the minimum. This prevents the "20 queries all pass admission simultaneously"
+/// burst — each new query arriving when RSS is elevated starts at minimum
+/// parallelism, limiting its hash table growth and total memory footprint.
 fn acquire_budget_inner(
     pool: &Arc<dyn MemoryPool>,
     avg_row_bytes: usize,
@@ -202,6 +209,22 @@ fn acquire_budget_inner(
     let mut target_partitions = configured_target_partitions.max(min_partitions);
     let mut batch_size = configured_batch_size.max(MIN_BATCH_SIZE);
 
+    // Proactive admission guard: only consult jemalloc RSS when the pool's own
+    // reservation accounting already shows pressure (>= admission threshold).
+    // This avoids the ~5µs jemalloc epoch.advance cost on the happy path.
+    if let Some(limit) = pool_limit(pool) {
+        let reserved = pool.reserved();
+        let threshold = crate::memory_guard::get_thresholds().admission;
+        let admission_bytes = (limit as f64 * threshold) as usize;
+        if reserved >= admission_bytes && crate::memory_guard::is_memory_pressured(limit) {
+            native_bridge_common::log_info!(
+                "Admission: pool reserved={}B (>={:.0}%) AND RSS at pressure — starting at min partitions={} batch_size={}",
+                reserved, threshold * 100.0, min_partitions, batch_size
+            );
+            target_partitions = min_partitions;
+        }
+    }
+
     loop {
         let phantom_bytes = compute_untracked_bytes_with_columns(
             target_partitions, batch_size, avg_row_bytes, num_columns,