From 0630ff07ae313c3820627cfdd24e95243e96b0e4 Mon Sep 17 00:00:00 2001
From: Brent Gardner <bgardner@squarelabs.net>
Date: Fri, 29 May 2026 12:17:13 -0600
Subject: [PATCH 01/17] add AccountingAllocator GlobalAlloc wrapper behind a
 feature flag
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Introduces a minimal `AccountingAllocator<A: GlobalAlloc>` that forwards
every alloc/realloc/dealloc unchanged to an inner allocator (default
`System`) and updates a thread-local `isize` byte counter on each call.
Reading the counter from another thread is allowed but only reflects
that thread's view — sum across threads for a global picture.

Gated by the new `memory-accounting` feature on `datafusion-sqllogictest`.
With the feature off the module is not compiled and the binary's
`#[global_allocator]` declaration is absent, so default builds are
completely unaffected. With the feature on the slt binary installs
the wrapper as its global allocator.

A unit test verifies the counter moves on a single-thread alloc/free
pair. Future commits will add the sampler and assertion plumbing on
top of this counter.
---
 datafusion/sqllogictest/Cargo.toml           |   4 +
 datafusion/sqllogictest/bin/sqllogictests.rs |   5 +
 datafusion/sqllogictest/src/accounting.rs    | 146 +++++++++++++++++++
 datafusion/sqllogictest/src/lib.rs           |   5 +
 4 files changed, 160 insertions(+)
 create mode 100644 datafusion/sqllogictest/src/accounting.rs

diff --git a/datafusion/sqllogictest/Cargo.toml b/datafusion/sqllogictest/Cargo.toml
index e2ffe1415a1fb..bbb288f82dfad 100644
--- a/datafusion/sqllogictest/Cargo.toml
+++ b/datafusion/sqllogictest/Cargo.toml
@@ -69,6 +69,10 @@ tokio-postgres = { version = "0.7.17", optional = true }
 [features]
 avro = ["datafusion/avro"]
 backtrace = ["datafusion/backtrace"]
+# Enable the `AccountingAllocator` `GlobalAlloc` wrapper and its thread-local
+# byte counter. The binary still has to declare `#[global_allocator]` for it
+# to actually take effect — building with this feature on alone is harmless.
+memory-accounting = []
 postgres = [
     "bytes",
     "chrono",
diff --git a/datafusion/sqllogictest/bin/sqllogictests.rs b/datafusion/sqllogictest/bin/sqllogictests.rs
index 69ae3a2fa7dd3..433dfefa94ace 100644
--- a/datafusion/sqllogictest/bin/sqllogictests.rs
+++ b/datafusion/sqllogictest/bin/sqllogictests.rs
@@ -15,6 +15,11 @@
 // specific language governing permissions and limitations
 // under the License.
 
+#[cfg(feature = "memory-accounting")]
+#[global_allocator]
+static GLOBAL: datafusion_sqllogictest::AccountingAllocator =
+    datafusion_sqllogictest::AccountingAllocator::system();
+
 use clap::{ColorChoice, Parser};
 use datafusion::common::instant::Instant;
 use datafusion::common::utils::get_available_parallelism;
diff --git a/datafusion/sqllogictest/src/accounting.rs b/datafusion/sqllogictest/src/accounting.rs
new file mode 100644
index 0000000000000..d84b88607169e
--- /dev/null
+++ b/datafusion/sqllogictest/src/accounting.rs
@@ -0,0 +1,146 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Allocation-counting `GlobalAlloc` wrapper.
+//!
+//! Wraps an inner allocator (default: `System`) and records every alloc/realloc/dealloc
+//! against a thread-local counter. Reads of the counter from another thread will
+//! not see this thread's pending state — sum across threads to get a global view.
+//!
+//! The counter can go negative on a thread that frees memory allocated elsewhere
+//! (e.g., `Arc` dropped on a worker that didn't construct it). The global sum
+//! across all threads is still correct.
+//!
+//! Compiled in only when the `memory-accounting` feature is on.
+//!
+//! # Usage
+//!
+//! ```ignore
+//! use datafusion_sqllogictest::AccountingAllocator;
+//!
+//! #[global_allocator]
+//! static GLOBAL: AccountingAllocator = AccountingAllocator::system();
+//!
+//! fn main() {
+//!     let before = datafusion_sqllogictest::bytes_in_use_this_thread();
+//!     let _v: Vec<u8> = vec![0; 1024];
+//!     let after = datafusion_sqllogictest::bytes_in_use_this_thread();
+//!     assert!(after - before >= 1024);
+//! }
+//! ```
+
+use std::alloc::{GlobalAlloc, Layout, System};
+use std::cell::Cell;
+
+thread_local! {
+    static BYTES_IN_USE: Cell<isize> = const { Cell::new(0) };
+}
+
+/// `GlobalAlloc` wrapper that counts bytes against a thread-local.
+///
+/// Forwards every operation unchanged to the inner allocator, so behavior
+/// (and performance) is identical aside from the atomic-free thread-local
+/// increment/decrement.
+pub struct AccountingAllocator<A: GlobalAlloc = System> {
+    inner: A,
+}
+
+impl<A: GlobalAlloc> AccountingAllocator<A> {
+    pub const fn new(inner: A) -> Self {
+        Self { inner }
+    }
+}
+
+impl AccountingAllocator<System> {
+    /// Convenience constructor for the typical `System`-backed case.
+    pub const fn system() -> Self {
+        Self { inner: System }
+    }
+}
+
+unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        // SAFETY: layout is forwarded unchanged.
+        let p = unsafe { self.inner.alloc(layout) };
+        if !p.is_null() {
+            // `try_with` so we don't panic if the TLS is being torn down.
+            let _ = BYTES_IN_USE.try_with(|c| c.set(c.get() + layout.size() as isize));
+        }
+        p
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        // SAFETY: caller upholds GlobalAlloc invariants; we forward unchanged.
+        unsafe { self.inner.dealloc(ptr, layout) };
+        let _ = BYTES_IN_USE.try_with(|c| c.set(c.get() - layout.size() as isize));
+    }
+
+    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+        // SAFETY: layout is forwarded unchanged.
+        let p = unsafe { self.inner.alloc_zeroed(layout) };
+        if !p.is_null() {
+            let _ = BYTES_IN_USE.try_with(|c| c.set(c.get() + layout.size() as isize));
+        }
+        p
+    }
+
+    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+        // SAFETY: caller upholds GlobalAlloc invariants; we forward unchanged.
+        let p = unsafe { self.inner.realloc(ptr, layout, new_size) };
+        if !p.is_null() {
+            let delta = new_size as isize - layout.size() as isize;
+            let _ = BYTES_IN_USE.try_with(|c| c.set(c.get() + delta));
+        }
+        p
+    }
+}
+
+/// Net bytes allocated minus deallocated on the current thread.
+///
+/// Cross-thread alloc/free will make individual thread values misleading;
+/// sum across threads (or use thread-pool snapshots) for a global picture.
+pub fn bytes_in_use_this_thread() -> isize {
+    BYTES_IN_USE.with(|c| c.get())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Verifies the counter moves in the expected direction on a single-thread
+    /// alloc and free pair. Requires this test binary to actually have the
+    /// allocator wired up via `#[global_allocator]` — handled below.
+    #[global_allocator]
+    static GLOBAL: AccountingAllocator = AccountingAllocator::system();
+
+    #[test]
+    fn counter_tracks_single_thread_alloc_and_free() {
+        let before = bytes_in_use_this_thread();
+        let v: Vec<u8> = vec![0u8; 8192];
+        let mid = bytes_in_use_this_thread();
+        assert!(
+            mid - before >= 8192,
+            "alloc not seen: before={before} mid={mid}"
+        );
+        drop(v);
+        let after = bytes_in_use_this_thread();
+        assert!(
+            mid - after >= 8192,
+            "free not seen: mid={mid} after={after}"
+        );
+    }
+}
diff --git a/datafusion/sqllogictest/src/lib.rs b/datafusion/sqllogictest/src/lib.rs
index 6b6c40365f855..4ea764647eebc 100644
--- a/datafusion/sqllogictest/src/lib.rs
+++ b/datafusion/sqllogictest/src/lib.rs
@@ -26,9 +26,14 @@
 
 //! DataFusion sqllogictest driver
 
+#[cfg(feature = "memory-accounting")]
+mod accounting;
 mod engines;
 mod test_file;
 
+#[cfg(feature = "memory-accounting")]
+pub use accounting::{AccountingAllocator, bytes_in_use_this_thread};
+
 pub use engines::CurrentlyExecutingSqlTracker;
 pub use engines::DFColumnType;
 pub use engines::DFOutput;

From ffa7277d5096a1a2f55ebe15100272c325ec3da7 Mon Sep 17 00:00:00 2001
From: Brent Gardner <bgardner@squarelabs.net>
Date: Fri, 29 May 2026 12:31:43 -0600
Subject: [PATCH 02/17] add global bank with threshold-based settlement and
 overdraft detection
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Extends AccountingAllocator from a per-thread counter into a budget model:

  * Allocations debit, deallocations credit.
  * Each thread maintains a local balance; when its magnitude crosses
    SETTLE_THRESHOLD (64 KB) it atomically transfers into a global bank
    and zeros out. Amortizes the atomic over ~thousands of allocs.
  * init_bank(budget) preloads the bank with a starting balance and clears
    the overdraft flag.
  * bank_balance() reads the bank with one relaxed load. Negative = the
    program has allocated more than init_bank seeded.
  * is_overdrawn() is a sticky flag set by the settling thread at the
    moment its fetch_add drops the bank past zero. Edge-triggered, so
    callers don't need a separate sampler to detect transitions.
  * clear_overdraft() resets the flag without touching the bank.
  * settle_thread_local() force-flushes the current thread's local into
    the bank — useful when a test wants an exact snapshot.

Hot-path cost is one branch + thread-local update; the atomic settle
fires roughly once per SETTLE_THRESHOLD bytes of net change per thread.
On a contended workload the bank's fetch_add runs ~1000 times/sec total
across all cores — effectively free compared to a per-op atomic that
would cost 5-10% of total runtime on hashagg-heavy code.

Four unit tests cover the new surface: debit/credit, init, overdraft
edge detection, and auto-settle on threshold crossing.
---
 datafusion/sqllogictest/src/accounting.rs | 239 +++++++++++++++++-----
 datafusion/sqllogictest/src/lib.rs        |   5 +-
 2 files changed, 190 insertions(+), 54 deletions(-)

diff --git a/datafusion/sqllogictest/src/accounting.rs b/datafusion/sqllogictest/src/accounting.rs
index d84b88607169e..bc12a653c6291 100644
--- a/datafusion/sqllogictest/src/accounting.rs
+++ b/datafusion/sqllogictest/src/accounting.rs
@@ -15,46 +15,127 @@
 // specific language governing permissions and limitations
 // under the License.
 
-//! Allocation-counting `GlobalAlloc` wrapper.
+//! Allocation-budget `GlobalAlloc` wrapper.
 //!
-//! Wraps an inner allocator (default: `System`) and records every alloc/realloc/dealloc
-//! against a thread-local counter. Reads of the counter from another thread will
-//! not see this thread's pending state — sum across threads to get a global view.
+//! The bank is a *budget*, not a counter. Allocations *debit* the bank,
+//! deallocations *credit* it. [`init_bank`] sets the starting budget.
+//! Going below zero is an overdraft — detectable via [`is_overdrawn`] or
+//! by reading [`bank_balance`] directly.
 //!
-//! The counter can go negative on a thread that frees memory allocated elsewhere
-//! (e.g., `Arc` dropped on a worker that didn't construct it). The global sum
-//! across all threads is still correct.
+//! Every `alloc`/`realloc`/`dealloc` updates a thread-local `isize` (also
+//! a balance — debit on alloc, credit on free). When that local's magnitude
+//! crosses [`SETTLE_THRESHOLD`] it transfers into the global bank with one
+//! atomic op and the local zeros out. This amortizes cache-line contention
+//! across thousands of allocations.
 //!
-//! Compiled in only when the `memory-accounting` feature is on.
-//!
-//! # Usage
-//!
-//! ```ignore
-//! use datafusion_sqllogictest::AccountingAllocator;
-//!
-//! #[global_allocator]
-//! static GLOBAL: AccountingAllocator = AccountingAllocator::system();
+//! [`bank_balance`] reads the global bank with a single relaxed atomic load.
+//! It lags reality by up to `SETTLE_THRESHOLD * num_threads` (≈1 MB on a
+//! 16-core box) — fine for catching multi-GB drift.
 //!
-//! fn main() {
-//!     let before = datafusion_sqllogictest::bytes_in_use_this_thread();
-//!     let _v: Vec<u8> = vec![0; 1024];
-//!     let after = datafusion_sqllogictest::bytes_in_use_this_thread();
-//!     assert!(after - before >= 1024);
-//! }
-//! ```
+//! Compiled in only when the `memory-accounting` feature is on.
 
 use std::alloc::{GlobalAlloc, Layout, System};
 use std::cell::Cell;
+use std::sync::atomic::{AtomicBool, AtomicIsize, Ordering};
+
+/// Net byte change at which a thread flushes its local count into the bank.
+/// 64 KB chosen to keep per-thread drift tight (≤1 MB on a 16-core box) while
+/// still settling rarely enough to make the bank's atomic op amortized-free.
+const SETTLE_THRESHOLD: isize = 64 * 1024;
+
+static BANK: AtomicIsize = AtomicIsize::new(0);
+
+/// Sticky flag: set the moment any thread's settle observes the bank crossing
+/// below zero. Cleared by [`init_bank`] or [`clear_overdraft`]. Reads are
+/// a single relaxed atomic load.
+static OVERDRAWN: AtomicBool = AtomicBool::new(false);
 
 thread_local! {
-    static BYTES_IN_USE: Cell<isize> = const { Cell::new(0) };
+    static LOCAL: Cell<isize> = const { Cell::new(0) };
+}
+
+/// Set the bank to `value` and clear the overdraft flag. Idempotent; safe
+/// to call from any thread.
+///
+/// Typical use: call once early in `main` to preload a memory budget before
+/// significant allocation occurs. Note that per-thread locals are NOT zeroed,
+/// so calling this mid-program leaves prior unsettled deltas in the locals
+/// that will eventually flush into the bank.
+pub fn init_bank(value: isize) {
+    BANK.store(value, Ordering::Relaxed);
+    OVERDRAWN.store(false, Ordering::Relaxed);
+}
+
+/// True iff any settle has ever seen the bank go below zero since the last
+/// [`init_bank`] / [`clear_overdraft`]. Single relaxed atomic load.
+pub fn is_overdrawn() -> bool {
+    OVERDRAWN.load(Ordering::Relaxed)
+}
+
+/// Clear the sticky overdraft flag without touching the bank.
+pub fn clear_overdraft() {
+    OVERDRAWN.store(false, Ordering::Relaxed);
+}
+
+/// Current bank balance. Allocs debit, deallocs credit, negative = overdraft.
+/// Single relaxed atomic load; cheap to call from any thread.
+///
+/// Lags reality by up to `SETTLE_THRESHOLD * num_threads`. To get an exact
+/// snapshot, drain each thread's local first via [`settle_thread_local`] —
+/// usually overkill.
+pub fn bank_balance() -> isize {
+    BANK.load(Ordering::Relaxed)
+}
+
+/// Current thread's local balance — not yet reflected in the global bank.
+/// Always in `(-SETTLE_THRESHOLD, +SETTLE_THRESHOLD)`. Sign matches the bank:
+/// negative on a thread that's net-allocated, positive on one that's net-freed.
+pub fn local_balance() -> isize {
+    LOCAL.with(|c| c.get())
+}
+
+/// Force the current thread to flush its local count into the bank.
+pub fn settle_thread_local() {
+    let _ = LOCAL.try_with(|c| {
+        let v = c.replace(0);
+        if v != 0 {
+            BANK.fetch_add(v, Ordering::Relaxed);
+        }
+    });
+}
+
+/// Record a delta against the current thread's local; flush to bank if we've
+/// crossed `±SETTLE_THRESHOLD`. Inlines into the alloc hot path.
+///
+/// Returns whether *this* settle observed the bank below zero. Callers in
+/// the allocator path ignore the result (we can't safely panic mid-alloc),
+/// but the function also sets the sticky [`OVERDRAWN`] flag so other code
+/// can detect the condition.
+#[inline(always)]
+fn track(delta: isize) -> bool {
+    let mut overdrew = false;
+    let _ = LOCAL.try_with(|c| {
+        let new = c.get() + delta;
+        if new >= SETTLE_THRESHOLD || new <= -SETTLE_THRESHOLD {
+            // `fetch_add` returns the OLD value. new_bank = old + new.
+            let old_bank = BANK.fetch_add(new, Ordering::Relaxed);
+            let new_bank = old_bank.wrapping_add(new);
+            if new_bank < 0 {
+                OVERDRAWN.store(true, Ordering::Relaxed);
+                overdrew = true;
+            }
+            c.set(0);
+        } else {
+            c.set(new);
+        }
+    });
+    overdrew
 }
 
-/// `GlobalAlloc` wrapper that counts bytes against a thread-local.
+/// `GlobalAlloc` wrapper that counts bytes against a thread-local + global bank.
 ///
-/// Forwards every operation unchanged to the inner allocator, so behavior
-/// (and performance) is identical aside from the atomic-free thread-local
-/// increment/decrement.
+/// Forwards every operation unchanged to the inner allocator; the bookkeeping
+/// is a thread-local update on the fast path plus an amortized atomic settle.
 pub struct AccountingAllocator<A: GlobalAlloc = System> {
     inner: A,
 }
@@ -77,8 +158,8 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> {
         // SAFETY: layout is forwarded unchanged.
         let p = unsafe { self.inner.alloc(layout) };
         if !p.is_null() {
-            // `try_with` so we don't panic if the TLS is being torn down.
-            let _ = BYTES_IN_USE.try_with(|c| c.set(c.get() + layout.size() as isize));
+            // Allocation debits the bank.
+            track(-(layout.size() as isize));
         }
         p
     }
@@ -86,14 +167,15 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> {
     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
         // SAFETY: caller upholds GlobalAlloc invariants; we forward unchanged.
         unsafe { self.inner.dealloc(ptr, layout) };
-        let _ = BYTES_IN_USE.try_with(|c| c.set(c.get() - layout.size() as isize));
+        // Free credits the bank.
+        track(layout.size() as isize);
     }
 
     unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
         // SAFETY: layout is forwarded unchanged.
         let p = unsafe { self.inner.alloc_zeroed(layout) };
         if !p.is_null() {
-            let _ = BYTES_IN_USE.try_with(|c| c.set(c.get() + layout.size() as isize));
+            track(-(layout.size() as isize));
         }
         p
     }
@@ -102,45 +184,96 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> {
         // SAFETY: caller upholds GlobalAlloc invariants; we forward unchanged.
         let p = unsafe { self.inner.realloc(ptr, layout, new_size) };
         if !p.is_null() {
-            let delta = new_size as isize - layout.size() as isize;
-            let _ = BYTES_IN_USE.try_with(|c| c.set(c.get() + delta));
+            // Growth debits, shrink credits.
+            track(layout.size() as isize - new_size as isize);
         }
         p
     }
 }
 
-/// Net bytes allocated minus deallocated on the current thread.
-///
-/// Cross-thread alloc/free will make individual thread values misleading;
-/// sum across threads (or use thread-pool snapshots) for a global picture.
-pub fn bytes_in_use_this_thread() -> isize {
-    BYTES_IN_USE.with(|c| c.get())
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
 
-    /// Verifies the counter moves in the expected direction on a single-thread
-    /// alloc and free pair. Requires this test binary to actually have the
-    /// allocator wired up via `#[global_allocator]` — handled below.
     #[global_allocator]
     static GLOBAL: AccountingAllocator = AccountingAllocator::system();
 
     #[test]
-    fn counter_tracks_single_thread_alloc_and_free() {
-        let before = bytes_in_use_this_thread();
+    fn alloc_debits_and_free_credits_bank() {
+        settle_thread_local();
+        let before = bank_balance();
+
         let v: Vec<u8> = vec![0u8; 8192];
-        let mid = bytes_in_use_this_thread();
+        settle_thread_local();
+        let mid = bank_balance();
+        // Alloc debited the bank → mid should be at least 8192 below before.
         assert!(
-            mid - before >= 8192,
-            "alloc not seen: before={before} mid={mid}"
+            before - mid >= 8192,
+            "alloc didn't debit: before={before} mid={mid}"
         );
+
         drop(v);
-        let after = bytes_in_use_this_thread();
+        settle_thread_local();
+        let after = bank_balance();
+        // Free credited the bank → after should be at least 8192 above mid.
+        assert!(
+            after - mid >= 8192,
+            "free didn't credit: mid={mid} after={after}"
+        );
+    }
+
+    #[test]
+    fn init_bank_sets_value() {
+        init_bank(1_000_000);
+        // The bank drifts from concurrent allocator activity in other tests
+        // (allocs debit), so we expect at-or-below the set value.
+        let v = bank_balance();
         assert!(
-            mid - after >= 8192,
-            "free not seen: mid={mid} after={after}"
+            (900_000..=1_000_000).contains(&v),
+            "init_bank didn't stick: v={v}"
         );
+        init_bank(0);
+    }
+
+    #[test]
+    fn overdraft_flag_trips_when_bank_goes_negative() {
+        init_bank(1024);
+        assert!(!is_overdrawn(), "flag should be clear after init_bank");
+
+        // Alloc large enough to cross SETTLE_THRESHOLD in one shot — settle
+        // drives bank from +1024 to roughly -(threshold+budget+overhead).
+        let _v: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 4096];
+        assert!(
+            is_overdrawn(),
+            "overdraft flag didn't trip; bank={}",
+            bank_balance()
+        );
+        assert!(
+            bank_balance() < 0,
+            "bank should be negative after overdraft; bank={}",
+            bank_balance()
+        );
+
+        clear_overdraft();
+        assert!(!is_overdrawn(), "clear_overdraft didn't clear the flag");
+
+        init_bank(0);
+    }
+
+    #[test]
+    fn threshold_settlement_flushes_to_bank() {
+        settle_thread_local();
+        let before_bank = bank_balance();
+
+        let v: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 1024];
+        // Crossing the threshold auto-settles; bank should have dropped by
+        // at least SETTLE_THRESHOLD without us calling settle_thread_local.
+        let bank_after_alloc = bank_balance();
+        assert!(
+            before_bank - bank_after_alloc >= SETTLE_THRESHOLD,
+            "bank didn't auto-settle on threshold crossing: \
+             before={before_bank} after_alloc={bank_after_alloc}"
+        );
+        drop(v);
     }
 }
diff --git a/datafusion/sqllogictest/src/lib.rs b/datafusion/sqllogictest/src/lib.rs
index 4ea764647eebc..aeaec2f329675 100644
--- a/datafusion/sqllogictest/src/lib.rs
+++ b/datafusion/sqllogictest/src/lib.rs
@@ -32,7 +32,10 @@ mod engines;
 mod test_file;
 
 #[cfg(feature = "memory-accounting")]
-pub use accounting::{AccountingAllocator, bytes_in_use_this_thread};
+pub use accounting::{
+    AccountingAllocator, bank_balance, clear_overdraft, init_bank, is_overdrawn,
+    local_balance, settle_thread_local,
+};
 
 pub use engines::CurrentlyExecutingSqlTracker;
 pub use engines::DFColumnType;

From a0a6be30e609ca7bbeaf86ec7210746e02890e3b Mon Sep 17 00:00:00 2001
From: Brent Gardner <bgardner@squarelabs.net>
Date: Fri, 29 May 2026 12:50:47 -0600
Subject: [PATCH 03/17] panic marked futures on overdraft via kill_on_overdraft
 scope

Adds the enforcement layer to the accounting allocator. Until now an
overdraft only set a sticky flag; the program kept running. This commit
gives the allocator the authority to panic the in-flight query work that
caused the overdraft -- but only work that has explicitly opted in.

The opt-in is a `kill_on_overdraft(future).await` scope built on
`tokio::task_local!`. While `f` is being polled, the very next debit
that observes the bank below zero panics with a typed `OverdraftPanic`
payload. Drops on the unwind do not re-panic because the same task-local
holds a `Cell<bool>` re-entry guard that flips to `true` the moment we
decide to panic.

Crucially, allocations on threads NOT inside such a scope (tokio reactor,
hyper, anything system) are untouched -- `KILL_GUARD.try_with(...)` returns
`Err` outside the scope and the alloc fast path skips the kill check
entirely. So the same worker thread can run a datafusion query (killable)
and a system task (exempt) back-to-back without leaking state, because
the marker lives on the future, not the thread.

Why a task_local rather than a thread_local + RAII guard: futures can
migrate between worker threads across `.await` boundaries. A thread_local
guard would mark one thread and clean up a different one, leaking the
mark. The task_local moves with the future and is unwound by tokio's own
machinery on completion / panic / cancellation -- no manual cleanup is
even possible.

Why the panic check fires only on debits: dealloc runs from Drop during
unwinding; a panic in Drop double-faults and aborts the process. Credit
operations therefore skip the kill check entirely -- they can't drive
the bank negative anyway.

Two new tests cover the enforcement surface: a wrapped future that
overdraws and is caught downcasting the OverdraftPanic payload, and an
identical workload outside the scope that overdraws without panicking.
All six accounting tests now serialize on a shared mutex since they
share the global bank.
---
 datafusion/sqllogictest/src/accounting.rs | 141 ++++++++++++++++++++--
 datafusion/sqllogictest/src/lib.rs        |   4 +-
 2 files changed, 133 insertions(+), 12 deletions(-)

diff --git a/datafusion/sqllogictest/src/accounting.rs b/datafusion/sqllogictest/src/accounting.rs
index bc12a653c6291..a20f10b408f52 100644
--- a/datafusion/sqllogictest/src/accounting.rs
+++ b/datafusion/sqllogictest/src/accounting.rs
@@ -32,10 +32,22 @@
 //! It lags reality by up to `SETTLE_THRESHOLD * num_threads` (≈1 MB on a
 //! 16-core box) — fine for catching multi-GB drift.
 //!
+//! # Enforcement
+//!
+//! By default the allocator only *counts* — overdrafts set the sticky
+//! [`is_overdrawn`] flag but allocations continue. Code that wants the
+//! allocator to actively kill its in-flight work wraps its futures in
+//! [`kill_on_overdraft`]. While such a future is being polled, the very
+//! next debit after the bank goes negative panics that future with an
+//! [`OverdraftPanic`] payload. Other futures (system tasks, HTTP handlers)
+//! polled on the same worker thread are unaffected — they're not inside
+//! the scope, so `track()` skips the check.
+//!
 //! Compiled in only when the `memory-accounting` feature is on.
 
 use std::alloc::{GlobalAlloc, Layout, System};
 use std::cell::Cell;
+use std::future::Future;
 use std::sync::atomic::{AtomicBool, AtomicIsize, Ordering};
 
 /// Net byte change at which a thread flushes its local count into the bank.
@@ -54,6 +66,41 @@ thread_local! {
     static LOCAL: Cell<isize> = const { Cell::new(0) };
 }
 
+tokio::task_local! {
+    /// Active iff the current task is inside a [`kill_on_overdraft`] scope.
+    /// Holding a `Cell<bool>` lets us track "we already started panicking"
+    /// without a separate thread-local, since the cell is task-scoped — it
+    /// disappears when the future unwinds past the scope, so the next task
+    /// on the same worker thread sees no enforcement.
+    static KILL_GUARD: Cell<bool>;
+}
+
+/// Payload attached to allocator-induced panics. Catch with:
+///
+/// ```ignore
+/// match std::panic::catch_unwind(|| { /* ... */ }) {
+///     Err(e) if e.is::<OverdraftPanic>() => { /* it was an overdraft */ }
+///     ...
+/// }
+/// ```
+#[derive(Debug, Clone)]
+pub struct OverdraftPanic {
+    /// Bank balance at the moment the panic fired (negative — that's the point).
+    pub bank_balance: isize,
+}
+
+/// Wrap a future so the allocator may panic it on overdraft.
+///
+/// While `f` is being polled, any allocation that finds the bank overdrawn
+/// causes an immediate `panic_any(OverdraftPanic)` on the polling thread.
+/// Drop runs of in-flight values during the unwind do not re-panic. After
+/// unwinding past the scope, the worker thread reverts to the default
+/// "count, but don't kill" mode, so any non-marked futures it polls next
+/// (system tasks, hyper handlers, etc.) are unaffected.
+pub async fn kill_on_overdraft<F: Future>(f: F) -> F::Output {
+    KILL_GUARD.scope(Cell::new(false), f).await
+}
+
 /// Set the bank to `value` and clear the overdraft flag. Idempotent; safe
 /// to call from any thread.
 ///
@@ -105,15 +152,10 @@ pub fn settle_thread_local() {
 }
 
 /// Record a delta against the current thread's local; flush to bank if we've
-/// crossed `±SETTLE_THRESHOLD`. Inlines into the alloc hot path.
-///
-/// Returns whether *this* settle observed the bank below zero. Callers in
-/// the allocator path ignore the result (we can't safely panic mid-alloc),
-/// but the function also sets the sticky [`OVERDRAWN`] flag so other code
-/// can detect the condition.
+/// crossed `±SETTLE_THRESHOLD`; on a debit, possibly kill the task. Inlines
+/// into the alloc hot path.
 #[inline(always)]
-fn track(delta: isize) -> bool {
-    let mut overdrew = false;
+fn track(delta: isize) {
     let _ = LOCAL.try_with(|c| {
         let new = c.get() + delta;
         if new >= SETTLE_THRESHOLD || new <= -SETTLE_THRESHOLD {
@@ -122,14 +164,36 @@ fn track(delta: isize) -> bool {
             let new_bank = old_bank.wrapping_add(new);
             if new_bank < 0 {
                 OVERDRAWN.store(true, Ordering::Relaxed);
-                overdrew = true;
             }
             c.set(0);
         } else {
             c.set(new);
         }
     });
-    overdrew
+
+    // Only debits can drive overdraft, and only debits run inside Drop-free
+    // code paths (alloc / alloc_zeroed / growing realloc). Credits happen
+    // during dealloc which runs inside Drop during unwinding — a panic there
+    // would double-fault and abort the process.
+    if delta < 0 && OVERDRAWN.load(Ordering::Relaxed) {
+        maybe_kill();
+    }
+}
+
+/// If the current task is inside a [`kill_on_overdraft`] scope and we haven't
+/// already started panicking for it, fire the kill panic.
+#[cold]
+#[inline(never)]
+fn maybe_kill() {
+    let _ = KILL_GUARD.try_with(|panicking| {
+        if !panicking.get() {
+            panicking.set(true);
+            let balance = BANK.load(Ordering::Relaxed);
+            std::panic::panic_any(OverdraftPanic {
+                bank_balance: balance,
+            });
+        }
+    });
 }
 
 /// `GlobalAlloc` wrapper that counts bytes against a thread-local + global bank.
@@ -194,12 +258,23 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use futures::FutureExt;
+    use std::sync::{Mutex, MutexGuard};
 
     #[global_allocator]
     static GLOBAL: AccountingAllocator = AccountingAllocator::system();
 
+    /// All accounting tests share the global bank, so they must run serially.
+    /// Each test takes this lock as its first line — a panicked test poisons
+    /// the lock; we recover the guard so subsequent tests still run.
+    static SERIAL: Mutex<()> = Mutex::new(());
+    fn serial() -> MutexGuard<'static, ()> {
+        SERIAL.lock().unwrap_or_else(|p| p.into_inner())
+    }
+
     #[test]
     fn alloc_debits_and_free_credits_bank() {
+        let _g = serial();
         settle_thread_local();
         let before = bank_balance();
 
@@ -224,6 +299,7 @@ mod tests {
 
     #[test]
     fn init_bank_sets_value() {
+        let _g = serial();
         init_bank(1_000_000);
         // The bank drifts from concurrent allocator activity in other tests
         // (allocs debit), so we expect at-or-below the set value.
@@ -237,6 +313,7 @@ mod tests {
 
     #[test]
     fn overdraft_flag_trips_when_bank_goes_negative() {
+        let _g = serial();
         init_bank(1024);
         assert!(!is_overdrawn(), "flag should be clear after init_bank");
 
@@ -260,8 +337,52 @@ mod tests {
         init_bank(0);
     }
 
+    #[tokio::test(flavor = "current_thread")]
+    async fn kill_on_overdraft_panics_inside_scope() {
+        let _g = serial();
+        init_bank(1024);
+
+        let result = std::panic::AssertUnwindSafe(kill_on_overdraft(async {
+            // First debit larger than SETTLE_THRESHOLD auto-settles and
+            // drives the bank to negative → next debit-with-overdraft kills.
+            // We do two big allocations so the second is the one that
+            // observes OVERDRAWN=true and panics.
+            let _v1: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 4096];
+            let _v2: Vec<u8> = vec![0u8; 1024];
+            unreachable!("should have panicked");
+        }))
+        .catch_unwind()
+        .await;
+
+        let payload = result.expect_err("future should have panicked");
+        let overdraft = payload
+            .downcast_ref::<OverdraftPanic>()
+            .expect("panic payload should be OverdraftPanic");
+        assert!(
+            overdraft.bank_balance < 0,
+            "payload should report negative balance; got {}",
+            overdraft.bank_balance
+        );
+
+        init_bank(0);
+    }
+
+    #[tokio::test(flavor = "current_thread")]
+    async fn no_kill_outside_scope() {
+        let _g = serial();
+        // No `kill_on_overdraft` wrapper -- overdraft should set the flag
+        // but allocations keep proceeding without panicking.
+        init_bank(1024);
+        let _v1: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 4096];
+        let _v2: Vec<u8> = vec![0u8; 1024];
+        assert!(is_overdrawn(), "should have overdrawn");
+        // Reaching here without panic is the whole point.
+        init_bank(0);
+    }
+
     #[test]
     fn threshold_settlement_flushes_to_bank() {
+        let _g = serial();
         settle_thread_local();
         let before_bank = bank_balance();
 
diff --git a/datafusion/sqllogictest/src/lib.rs b/datafusion/sqllogictest/src/lib.rs
index aeaec2f329675..cf58851931e58 100644
--- a/datafusion/sqllogictest/src/lib.rs
+++ b/datafusion/sqllogictest/src/lib.rs
@@ -33,8 +33,8 @@ mod test_file;
 
 #[cfg(feature = "memory-accounting")]
 pub use accounting::{
-    AccountingAllocator, bank_balance, clear_overdraft, init_bank, is_overdrawn,
-    local_balance, settle_thread_local,
+    AccountingAllocator, OverdraftPanic, bank_balance, clear_overdraft, init_bank,
+    is_overdrawn, kill_on_overdraft, local_balance, settle_thread_local,
 };
 
 pub use engines::CurrentlyExecutingSqlTracker;

From ba2eb393f05e3f920f9d27e9333423dfa8a6a624 Mon Sep 17 00:00:00 2001
From: Brent Gardner <bgardner@squarelabs.net>
Date: Fri, 29 May 2026 12:57:12 -0600
Subject: [PATCH 04/17] add bytes_used() so panic handlers can report actual
 allocation

Stores INITIAL_BUDGET alongside the bank on init_bank() so reporters can
express usage as a positive "bytes allocated since the budget was set"
number rather than a deeply-negative bank balance.

Needed by the upcoming SLT integration commit: when a query is killed by
OverdraftPanic, the handler logs "actual N MB / supposed M MB" where
"actual" comes from bytes_used() and "supposed" comes from DataFusion's
voluntary MemoryPool.reserved(). The delta is the accounting bug.
---
 datafusion/sqllogictest/src/accounting.rs | 17 +++++++++++++++--
 datafusion/sqllogictest/src/lib.rs        |  4 ++--
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/datafusion/sqllogictest/src/accounting.rs b/datafusion/sqllogictest/src/accounting.rs
index a20f10b408f52..41957cb21bab3 100644
--- a/datafusion/sqllogictest/src/accounting.rs
+++ b/datafusion/sqllogictest/src/accounting.rs
@@ -57,6 +57,11 @@ const SETTLE_THRESHOLD: isize = 64 * 1024;
 
 static BANK: AtomicIsize = AtomicIsize::new(0);
 
+/// Last value passed to [`init_bank`]. Used by [`bytes_used`] to express
+/// "how much has been allocated since the budget was set" as a positive
+/// number, which is more readable than a deeply-negative balance.
+static INITIAL_BUDGET: AtomicIsize = AtomicIsize::new(0);
+
 /// Sticky flag: set the moment any thread's settle observes the bank crossing
 /// below zero. Cleared by [`init_bank`] or [`clear_overdraft`]. Reads are
 /// a single relaxed atomic load.
@@ -101,18 +106,26 @@ pub async fn kill_on_overdraft<F: Future>(f: F) -> F::Output {
     KILL_GUARD.scope(Cell::new(false), f).await
 }
 
-/// Set the bank to `value` and clear the overdraft flag. Idempotent; safe
-/// to call from any thread.
+/// Set the bank to `value`, remember it as the budget for [`bytes_used`],
+/// and clear the overdraft flag. Idempotent; safe to call from any thread.
 ///
 /// Typical use: call once early in `main` to preload a memory budget before
 /// significant allocation occurs. Note that per-thread locals are NOT zeroed,
 /// so calling this mid-program leaves prior unsettled deltas in the locals
 /// that will eventually flush into the bank.
 pub fn init_bank(value: isize) {
+    INITIAL_BUDGET.store(value, Ordering::Relaxed);
     BANK.store(value, Ordering::Relaxed);
     OVERDRAWN.store(false, Ordering::Relaxed);
 }
 
+/// Bytes allocated since the last [`init_bank`] — i.e., `budget - bank_balance`.
+/// Always positive (or zero) under normal conditions; only meaningful relative
+/// to the most recent `init_bank` call.
+pub fn bytes_used() -> isize {
+    INITIAL_BUDGET.load(Ordering::Relaxed) - BANK.load(Ordering::Relaxed)
+}
+
 /// True iff any settle has ever seen the bank go below zero since the last
 /// [`init_bank`] / [`clear_overdraft`]. Single relaxed atomic load.
 pub fn is_overdrawn() -> bool {
diff --git a/datafusion/sqllogictest/src/lib.rs b/datafusion/sqllogictest/src/lib.rs
index cf58851931e58..5ad60dce6e6c0 100644
--- a/datafusion/sqllogictest/src/lib.rs
+++ b/datafusion/sqllogictest/src/lib.rs
@@ -33,8 +33,8 @@ mod test_file;
 
 #[cfg(feature = "memory-accounting")]
 pub use accounting::{
-    AccountingAllocator, OverdraftPanic, bank_balance, clear_overdraft, init_bank,
-    is_overdrawn, kill_on_overdraft, local_balance, settle_thread_local,
+    AccountingAllocator, OverdraftPanic, bank_balance, bytes_used, clear_overdraft,
+    init_bank, is_overdrawn, kill_on_overdraft, local_balance, settle_thread_local,
 };
 
 pub use engines::CurrentlyExecutingSqlTracker;

From 0c469497f423f2d9e3e100e41b0c9218dea9621a Mon Sep 17 00:00:00 2001
From: Brent Gardner <bgardner@squarelabs.net>
Date: Fri, 29 May 2026 13:05:07 -0600
Subject: [PATCH 05/17] Hook it all up

---
 datafusion/sqllogictest/bin/sqllogictests.rs  | 34 +++++++++++
 datafusion/sqllogictest/src/accounting.rs     | 18 +++++-
 .../src/engines/datafusion_engine/runner.rs   | 59 ++++++++++++++++++-
 datafusion/sqllogictest/src/lib.rs            |  3 +-
 datafusion/sqllogictest/src/test_context.rs   | 23 +++++++-
 5 files changed, 133 insertions(+), 4 deletions(-)

diff --git a/datafusion/sqllogictest/bin/sqllogictests.rs b/datafusion/sqllogictest/bin/sqllogictests.rs
index 433dfefa94ace..26b945eee7cb4 100644
--- a/datafusion/sqllogictest/bin/sqllogictests.rs
+++ b/datafusion/sqllogictest/bin/sqllogictests.rs
@@ -142,6 +142,22 @@ async fn run_tests() -> Result<()> {
 
     options.warn_on_ignored();
 
+    #[cfg(feature = "memory-accounting")]
+    if let Some(mb) = options.memory_tracker_mb {
+        let tracker_bytes = mb.saturating_mul(1024 * 1024);
+        let killer_bytes = (tracker_bytes as f64
+            * (1.0 + options.memory_killer_overhead_pct as f64 / 100.0))
+            as isize;
+        datafusion_sqllogictest::set_memory_tracker_limit(tracker_bytes);
+        datafusion_sqllogictest::init_bank(killer_bytes);
+        log::info!(
+            "memory-accounting on: tracker={mb} MB, killer={} MB \
+             (overhead +{}%)",
+            killer_bytes / (1024 * 1024),
+            options.memory_killer_overhead_pct
+        );
+    }
+
     // Print parallelism info for debugging CI performance
     eprintln!(
         "Running with {} test threads (available parallelism: {})",
@@ -915,6 +931,24 @@ struct Options {
         default_value_t = ColorChoice::Auto
     )]
     color: ColorChoice,
+
+    #[clap(
+        long,
+        help = "Cap DataFusion's voluntary MemoryPool at this many MB. Requires \
+                the memory-accounting feature; ignored without it. 0 = unbounded \
+                (current default)."
+    )]
+    memory_tracker_mb: Option<usize>,
+
+    #[clap(
+        long,
+        default_value_t = 25,
+        help = "Allocator kill threshold expressed as a percentage above \
+                --memory-tracker-mb. When the real bytes allocated since query \
+                start exceed memory_tracker_mb * (1 + this/100), the in-flight \
+                query is panicked via OverdraftPanic and logged as a regression."
+    )]
+    memory_killer_overhead_pct: u64,
 }
 
 impl Options {
diff --git a/datafusion/sqllogictest/src/accounting.rs b/datafusion/sqllogictest/src/accounting.rs
index 41957cb21bab3..ddb081a10ee49 100644
--- a/datafusion/sqllogictest/src/accounting.rs
+++ b/datafusion/sqllogictest/src/accounting.rs
@@ -48,7 +48,7 @@
 use std::alloc::{GlobalAlloc, Layout, System};
 use std::cell::Cell;
 use std::future::Future;
-use std::sync::atomic::{AtomicBool, AtomicIsize, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicIsize, AtomicUsize, Ordering};
 
 /// Net byte change at which a thread flushes its local count into the bank.
 /// 64 KB chosen to keep per-thread drift tight (≤1 MB on a 16-core box) while
@@ -126,6 +126,22 @@ pub fn bytes_used() -> isize {
     INITIAL_BUDGET.load(Ordering::Relaxed) - BANK.load(Ordering::Relaxed)
 }
 
+/// Cross-module config for DataFusion's voluntary `MemoryPool` limit, set
+/// from the SLT binary's CLI and read by test_context when building each
+/// per-file `RuntimeEnv`. Zero means "use the default `UnboundedMemoryPool`".
+static MEMORY_TRACKER_LIMIT: AtomicUsize = AtomicUsize::new(0);
+
+/// Set the size (in bytes) the per-file `MemoryPool` should be built with.
+/// Zero (the default) leaves the existing `UnboundedMemoryPool` behavior.
+pub fn set_memory_tracker_limit(bytes: usize) {
+    MEMORY_TRACKER_LIMIT.store(bytes, Ordering::Relaxed);
+}
+
+/// Current `MemoryPool` limit configured via [`set_memory_tracker_limit`].
+pub fn memory_tracker_limit() -> usize {
+    MEMORY_TRACKER_LIMIT.load(Ordering::Relaxed)
+}
+
 /// True iff any settle has ever seen the bank go below zero since the last
 /// [`init_bank`] / [`clear_overdraft`]. Single relaxed atomic load.
 pub fn is_overdrawn() -> bool {
diff --git a/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs b/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs
index 08facc48005dc..a20cdc4300781 100644
--- a/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs
+++ b/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs
@@ -83,6 +83,63 @@ impl DataFusion {
         self
     }
 
+    /// Run a single query through the engine, wrapped (when built with the
+    /// `memory-accounting` feature) in a kill-on-overdraft scope that turns
+    /// allocator-detected budget overruns into clean `Err` results instead
+    /// of process crashes.
+    async fn run_one(&self, sql: &str) -> Result<DFOutput> {
+        #[cfg(feature = "memory-accounting")]
+        {
+            use crate::{OverdraftPanic, init_bank, kill_on_overdraft};
+            use futures::FutureExt;
+
+            // Reset the bank before each statement so per-query budgets are
+            // independent. The post-run drop chain credits everything back
+            // anyway, but this also clears the sticky overdraft flag in case
+            // a prior statement tripped it.
+            let budget = crate::bank_balance().max(crate::memory_tracker_limit() as isize);
+            init_bank(budget);
+
+            let fut = kill_on_overdraft(run_query(
+                &self.ctx,
+                is_spark_path(&self.relative_path),
+                sql,
+            ));
+
+            return match std::panic::AssertUnwindSafe(fut).catch_unwind().await {
+                Ok(r) => r,
+                Err(payload) => {
+                    if let Some(od) = payload.downcast_ref::<OverdraftPanic>() {
+                        let actual = (crate::bytes_used().max(0) as u64) / (1024 * 1024);
+                        let supposed = (self.ctx.runtime_env().memory_pool.reserved()
+                            as u64)
+                            / (1024 * 1024);
+                        warn!(
+                            "[{}] killed by allocator overdraft: \
+                             actual={actual} MB, supposed={supposed} MB \
+                             (bank balance at panic = {} bytes); sql={sql:?}",
+                            self.relative_path.display(),
+                            od.bank_balance
+                        );
+                        // Reset bank to the configured budget so the next
+                        // statement in this file gets a fresh budget.
+                        init_bank(crate::memory_tracker_limit() as isize);
+                        Err(DFSqlLogicTestError::Other(format!(
+                            "allocator overdraft: actual={actual} MB supposed={supposed} MB"
+                        )))
+                    } else {
+                        // Not our panic — re-raise so test runner sees it.
+                        std::panic::resume_unwind(payload);
+                    }
+                }
+            };
+        }
+        #[cfg(not(feature = "memory-accounting"))]
+        {
+            run_query(&self.ctx, is_spark_path(&self.relative_path), sql).await
+        }
+    }
+
     fn update_slow_count(&self) {
         let msg = self.pb.message();
         let split: Vec<&str> = msg.split(" ").collect();
@@ -154,7 +211,7 @@ impl sqllogictest::AsyncDB for DataFusion {
         let tracked_sql = self.currently_executing_sql_tracker.set_sql(sql);
 
         let start = Instant::now();
-        let result = run_query(&self.ctx, is_spark_path(&self.relative_path), sql).await;
+        let result = self.run_one(sql).await;
         let duration = start.elapsed();
 
         self.currently_executing_sql_tracker.remove_sql(tracked_sql);
diff --git a/datafusion/sqllogictest/src/lib.rs b/datafusion/sqllogictest/src/lib.rs
index 5ad60dce6e6c0..e4b3f09693d35 100644
--- a/datafusion/sqllogictest/src/lib.rs
+++ b/datafusion/sqllogictest/src/lib.rs
@@ -34,7 +34,8 @@ mod test_file;
 #[cfg(feature = "memory-accounting")]
 pub use accounting::{
     AccountingAllocator, OverdraftPanic, bank_balance, bytes_used, clear_overdraft,
-    init_bank, is_overdrawn, kill_on_overdraft, local_balance, settle_thread_local,
+    init_bank, is_overdrawn, kill_on_overdraft, local_balance, memory_tracker_limit,
+    set_memory_tracker_limit, settle_thread_local,
 };
 
 pub use engines::CurrentlyExecutingSqlTracker;
diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs
index 0edde71b939f4..590db6ddfa429 100644
--- a/datafusion/sqllogictest/src/test_context.rs
+++ b/datafusion/sqllogictest/src/test_context.rs
@@ -57,6 +57,8 @@ use async_trait::async_trait;
 use datafusion::common::cast::as_float64_array;
 use datafusion::execution::SessionStateBuilder;
 use datafusion::execution::runtime_env::RuntimeEnv;
+#[cfg(feature = "memory-accounting")]
+use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use log::info;
 use sqlparser::ast;
 use tempfile::TempDir;
@@ -86,6 +88,25 @@ impl TypePlanner for SqlLogicTestTypePlanner {
     }
 }
 
+/// Construct the per-file `RuntimeEnv`. With the `memory-accounting` feature
+/// on and a non-zero `memory_tracker_limit()` configured, this builds a
+/// finite-pool runtime so DataFusion's voluntary tracker can be observed
+/// against the allocator's ground-truth bank. Otherwise falls back to the
+/// historical `RuntimeEnv::default()` (unbounded pool).
+fn build_runtime_env() -> RuntimeEnv {
+    #[cfg(feature = "memory-accounting")]
+    {
+        let limit = crate::memory_tracker_limit();
+        if limit > 0 {
+            return RuntimeEnvBuilder::new()
+                .with_memory_limit(limit, 1.0)
+                .build()
+                .expect("RuntimeEnvBuilder::build with memory limit");
+        }
+    }
+    RuntimeEnv::default()
+}
+
 impl TestContext {
     pub fn new(ctx: SessionContext) -> Self {
         Self {
@@ -103,7 +124,7 @@ impl TestContext {
         let config = SessionConfig::new()
             // hardcode target partitions so plans are deterministic
             .with_target_partitions(4);
-        let runtime = Arc::new(RuntimeEnv::default());
+        let runtime = Arc::new(build_runtime_env());
 
         let mut state_builder = SessionStateBuilder::new()
             .with_config(config)

From 0cfe44c7f9ec1395829140290c0f06040fabd1a8 Mon Sep 17 00:00:00 2001
From: Brent Gardner <bgardner@squarelabs.net>
Date: Fri, 29 May 2026 13:20:08 -0600
Subject: [PATCH 06/17] factor reset_bank() out of the runner

The per-statement reset in run_one was doing
  let budget = bank_balance().max(memory_tracker_limit() as isize);
  init_bank(budget);
which was clever-for-the-wrong-reasons: it tried to pick the larger of
the current drifted bank and the configured tracker limit, conflating
"reset" with "init" with "configure". And the post-panic reset path used
init_bank(memory_tracker_limit() as isize) which dropped the +25% killer
overhead each time the bank reset (so the first overdraft permanently
shrank the budget for every subsequent statement in the file).

Replace both with reset_bank(), which restores BANK to whatever value
was passed to init_bank() and clears the overdraft flag. INITIAL_BUDGET
remembers that value across the whole run, so per-statement resets land
on the same number main set during startup.
---
 datafusion/sqllogictest/src/accounting.rs      | 12 ++++++++++++
 .../src/engines/datafusion_engine/runner.rs    | 18 ++++++++----------
 datafusion/sqllogictest/src/lib.rs             |  2 +-
 3 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/datafusion/sqllogictest/src/accounting.rs b/datafusion/sqllogictest/src/accounting.rs
index ddb081a10ee49..abb3318f823cf 100644
--- a/datafusion/sqllogictest/src/accounting.rs
+++ b/datafusion/sqllogictest/src/accounting.rs
@@ -126,6 +126,18 @@ pub fn bytes_used() -> isize {
     INITIAL_BUDGET.load(Ordering::Relaxed) - BANK.load(Ordering::Relaxed)
 }
 
+/// Reset the bank to the budget last passed to [`init_bank`] and clear the
+/// overdraft flag. No-op if `init_bank` has never been called (budget = 0).
+///
+/// Intended to be called between independent units of work — e.g., between
+/// SQL statements — so each gets a fresh budget regardless of whether the
+/// previous one drifted the bank or tripped overdraft.
+pub fn reset_bank() {
+    let budget = INITIAL_BUDGET.load(Ordering::Relaxed);
+    BANK.store(budget, Ordering::Relaxed);
+    OVERDRAWN.store(false, Ordering::Relaxed);
+}
+
 /// Cross-module config for DataFusion's voluntary `MemoryPool` limit, set
 /// from the SLT binary's CLI and read by test_context when building each
 /// per-file `RuntimeEnv`. Zero means "use the default `UnboundedMemoryPool`".
diff --git a/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs b/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs
index a20cdc4300781..5e35b9c981e40 100644
--- a/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs
+++ b/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs
@@ -90,15 +90,13 @@ impl DataFusion {
     async fn run_one(&self, sql: &str) -> Result<DFOutput> {
         #[cfg(feature = "memory-accounting")]
         {
-            use crate::{OverdraftPanic, init_bank, kill_on_overdraft};
+            use crate::{OverdraftPanic, kill_on_overdraft, reset_bank};
             use futures::FutureExt;
 
-            // Reset the bank before each statement so per-query budgets are
-            // independent. The post-run drop chain credits everything back
-            // anyway, but this also clears the sticky overdraft flag in case
-            // a prior statement tripped it.
-            let budget = crate::bank_balance().max(crate::memory_tracker_limit() as isize);
-            init_bank(budget);
+            // Fresh budget for this statement. Independent of whatever the
+            // previous statement did (or panicked through). Also clears the
+            // sticky overdraft flag.
+            reset_bank();
 
             let fut = kill_on_overdraft(run_query(
                 &self.ctx,
@@ -121,9 +119,9 @@ impl DataFusion {
                             self.relative_path.display(),
                             od.bank_balance
                         );
-                        // Reset bank to the configured budget so the next
-                        // statement in this file gets a fresh budget.
-                        init_bank(crate::memory_tracker_limit() as isize);
+                        // Pre-emptive reset so the next statement starts clean
+                        // even if some drop chain didn't fully credit the bank.
+                        reset_bank();
                         Err(DFSqlLogicTestError::Other(format!(
                             "allocator overdraft: actual={actual} MB supposed={supposed} MB"
                         )))
diff --git a/datafusion/sqllogictest/src/lib.rs b/datafusion/sqllogictest/src/lib.rs
index e4b3f09693d35..5d206f5bf7b76 100644
--- a/datafusion/sqllogictest/src/lib.rs
+++ b/datafusion/sqllogictest/src/lib.rs
@@ -35,7 +35,7 @@ mod test_file;
 pub use accounting::{
     AccountingAllocator, OverdraftPanic, bank_balance, bytes_used, clear_overdraft,
     init_bank, is_overdrawn, kill_on_overdraft, local_balance, memory_tracker_limit,
-    set_memory_tracker_limit, settle_thread_local,
+    reset_bank, set_memory_tracker_limit, settle_thread_local,
 };
 
 pub use engines::CurrentlyExecutingSqlTracker;

From 7d0bc52921029babce34ceb555b45320019627f7 Mon Sep 17 00:00:00 2001
From: Brent Gardner <bgardner@squarelabs.net>
Date: Fri, 29 May 2026 13:26:19 -0600
Subject: [PATCH 07/17] clippy + fmt: silence await-holding-lock on serialized
 async test

The kill_on_overdraft_panics_inside_scope test holds the SERIAL Mutex
guard across `.await` to serialize against the shared global bank.
Clippy's await-holding-lock lint flags this since in general holding a
sync Mutex across await can deadlock other tokio tasks, but here the
test uses `flavor = "current_thread"` so there's no concurrent task to
block, and the guard exists precisely to keep other accounting tests
from concurrently mutating the bank for the duration of the body.

Annotate with #[expect(...)] (with a reason) rather than #[allow(...)]
to satisfy clippy::allow_attributes; if the lint stops firing in the
future the expect will turn into an error and we'll revisit.
---
 datafusion/sqllogictest/src/accounting.rs | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/datafusion/sqllogictest/src/accounting.rs b/datafusion/sqllogictest/src/accounting.rs
index abb3318f823cf..49c4236602fba 100644
--- a/datafusion/sqllogictest/src/accounting.rs
+++ b/datafusion/sqllogictest/src/accounting.rs
@@ -378,6 +378,13 @@ mod tests {
         init_bank(0);
     }
 
+    // Holding the std Mutex guard across `.await` is fine here: the test uses
+    // a current-thread runtime so there's no other task to block, and we need
+    // exclusive access to the global bank for the entire test body.
+    #[expect(
+        clippy::await_holding_lock,
+        reason = "test uses current-thread runtime; the guard exists to serialize tests against the shared global bank"
+    )]
     #[tokio::test(flavor = "current_thread")]
     async fn kill_on_overdraft_panics_inside_scope() {
         let _g = serial();

From a6258397afed847710e1647c2b3d086fdb0bfee2 Mon Sep 17 00:00:00 2001
From: Brent Gardner <bgardner@squarelabs.net>
Date: Fri, 29 May 2026 13:31:55 -0600
Subject: [PATCH 08/17] reframe CLI as total-memory + datafusion-fraction

Old shape inverted DataFusion's intended `with_memory_limit(max, fraction)`
contract:

  --memory-tracker-mb 1024
  --memory-killer-overhead-pct 25
    -> DF pool sized to 1024 MB
    -> bank set to 1280 MB

That pretended DF gets 100% of some smaller pie and bolted "overhead"
onto the top, instead of starting from the actual total budget. With
the allocator bank tracking ALL process allocations (DF, hyper, tokio,
gRPC, gossip, etc.), the natural framing is:

  --total-memory-mb 1024
  --datafusion-memory-fraction 0.6
    -> bank set to 1024 MB (the whole pie)
    -> DF pool sized to 614 MB (its share)
    -> remaining 410 MB is non-DF budget

The bank trips when total RSS exceeds the OS-visible limit, regardless
of which consumer caused the drift. The DF pool trips when DF itself
exceeds its voluntary share. Both reads are meaningful and the relation
between them is now explicit instead of fudged.

CLI:
  --total-memory-mb (was --memory-tracker-mb)
  --datafusion-memory-fraction f64 default 0.6 (replaces --memory-killer-overhead-pct)

Startup log now prints all three numbers so it's obvious what each
knob is doing:
  memory-accounting on: total=1024 MB, datafusion=614 MB (fraction 0.60),
                        non-datafusion budget=409 MB
---
 datafusion/sqllogictest/bin/sqllogictests.rs | 49 +++++++++++---------
 1 file changed, 28 insertions(+), 21 deletions(-)

diff --git a/datafusion/sqllogictest/bin/sqllogictests.rs b/datafusion/sqllogictest/bin/sqllogictests.rs
index 26b945eee7cb4..92fb2351df957 100644
--- a/datafusion/sqllogictest/bin/sqllogictests.rs
+++ b/datafusion/sqllogictest/bin/sqllogictests.rs
@@ -143,18 +143,22 @@ async fn run_tests() -> Result<()> {
     options.warn_on_ignored();
 
     #[cfg(feature = "memory-accounting")]
-    if let Some(mb) = options.memory_tracker_mb {
-        let tracker_bytes = mb.saturating_mul(1024 * 1024);
-        let killer_bytes = (tracker_bytes as f64
-            * (1.0 + options.memory_killer_overhead_pct as f64 / 100.0))
-            as isize;
-        datafusion_sqllogictest::set_memory_tracker_limit(tracker_bytes);
-        datafusion_sqllogictest::init_bank(killer_bytes);
+    if let Some(total_mb) = options.total_memory_mb {
+        let total_bytes = total_mb.saturating_mul(1024 * 1024);
+        let df_bytes = (total_bytes as f64 * options.datafusion_memory_fraction) as usize;
+        // The DataFusion MemoryPool gets its slice of the total budget; the
+        // allocator bank tracks the whole pie (DF + everything else in the
+        // process). A kill_on_overdraft-marked future panics when total
+        // process allocations exceed `total_bytes`, regardless of whether
+        // DF's voluntary tracker noticed.
+        datafusion_sqllogictest::set_memory_tracker_limit(df_bytes);
+        datafusion_sqllogictest::init_bank(total_bytes as isize);
         log::info!(
-            "memory-accounting on: tracker={mb} MB, killer={} MB \
-             (overhead +{}%)",
-            killer_bytes / (1024 * 1024),
-            options.memory_killer_overhead_pct
+            "memory-accounting on: total={total_mb} MB, datafusion={} MB \
+             (fraction {:.2}), non-datafusion budget={} MB",
+            df_bytes / (1024 * 1024),
+            options.datafusion_memory_fraction,
+            (total_bytes - df_bytes) / (1024 * 1024)
         );
     }
 
@@ -934,21 +938,24 @@ struct Options {
 
     #[clap(
         long,
-        help = "Cap DataFusion's voluntary MemoryPool at this many MB. Requires \
-                the memory-accounting feature; ignored without it. 0 = unbounded \
-                (current default)."
+        help = "Total memory budget for the whole process in MB (cgroup limit, \
+                container size, etc.). The allocator bank is set to this value, \
+                and panics any kill_on_overdraft-marked future as soon as real \
+                bytes allocated since startup exceed it. Requires the \
+                memory-accounting feature; ignored without it."
     )]
-    memory_tracker_mb: Option<usize>,
+    total_memory_mb: Option<usize>,
 
     #[clap(
         long,
-        default_value_t = 25,
-        help = "Allocator kill threshold expressed as a percentage above \
-                --memory-tracker-mb. When the real bytes allocated since query \
-                start exceed memory_tracker_mb * (1 + this/100), the in-flight \
-                query is panicked via OverdraftPanic and logged as a regression."
+        default_value_t = 0.6,
+        help = "DataFusion's share of --total-memory-mb. The per-file \
+                MemoryPool is sized at total * this fraction; the rest of the \
+                process budget covers hyper, tokio internals, gRPC, and other \
+                non-DF consumers. The bank catches violations of the overall \
+                budget regardless of which consumer caused the drift."
     )]
-    memory_killer_overhead_pct: u64,
+    datafusion_memory_fraction: f64,
 }
 
 impl Options {

From 996954d6f92ce4dd1691cb62b2633d7092990dbb Mon Sep 17 00:00:00 2001
From: Brent Gardner <bgardner@squarelabs.net>
Date: Fri, 29 May 2026 14:31:34 -0600
Subject: [PATCH 09/17] turn on memory-accounting in the main SLT CI job with a
 loose budget

Wires the new feature into the existing full-corpus SLT run inside
the `verify-benchmark-results` job (which despite its name runs every
.slt file with `INCLUDE_TPCH=true`). Budget is intentionally loose:
16 GB total, 60% to DataFusion, ~6.4 GB to non-DF consumers. The
existing test corpus peaks well under 1 GB, so nothing should fail.

What this validates on every PR:
  - `--features memory-accounting` still compiles cleanly
  - AccountingAllocator is installed as #[global_allocator]
  - init_bank, reset_bank, and bytes_used run on every statement
  - Each query goes through kill_on_overdraft + catch_unwind without
    introducing measurable overhead or panicking incorrectly
  - The per-file MemoryPool gets sized via build_runtime_env

If the integration bit-rots in any future PR, CI catches it. Once
we have confidence in the wiring we can tighten the budget in-place
or graduate this into a dedicated job that exercises tight budgets
to surface actual memory-accounting bugs in DF operators.
---
 .github/workflows/rust.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml
index 5af7dc418c8d9..ac23a9775b8dd 100644
--- a/.github/workflows/rust.yml
+++ b/.github/workflows/rust.yml
@@ -471,7 +471,7 @@ jobs:
           export RUST_MIN_STACK=20971520
           export TPCH_DATA=`realpath datafusion/sqllogictest/test_files/tpch/data`
           cargo test plan_q --package datafusion-benchmarks --profile ci --features=ci -- --test-threads=1
-          INCLUDE_TPCH=true cargo test --features backtrace,parquet_encryption,substrait --profile ci --package datafusion-sqllogictest --test sqllogictests
+          INCLUDE_TPCH=true cargo test --features backtrace,parquet_encryption,substrait,memory-accounting --profile ci --package datafusion-sqllogictest --test sqllogictests -- --total-memory-mb 16384 --datafusion-memory-fraction 0.6
       - name: Verify Working Directory Clean
         run: git diff --exit-code
 

From b3ab387546297018f134b1bcdac4d6c4467f4e94 Mon Sep 17 00:00:00 2001
From: Brent Gardner <bgardner@squarelabs.net>
Date: Mon, 1 Jun 2026 10:05:52 -0600
Subject: [PATCH 10/17] per-file Tokio runtime in SLT runner with thread-local
 context-id
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Each SLT file now runs in its own multi-thread Tokio runtime hosted via
spawn_blocking on the outer orchestration runtime. Every worker thread
of a per-file runtime gets stamped with a unique context-id via
on_thread_start. Nothing reads the id yet — this is plumbing for a
follow-up that keys the bank balance off it so SET datafusion.runtime.
memory_limit can retune a per-context budget instead of fighting the
global one.

Why per-file runtimes:
  - The shared multi-thread runtime work-steals across files, so an
    OS thread serves many SessionContexts over its lifetime. A
    thread-local context-id read from inside the allocator hook would
    point at whichever file's task most recently set it — useless.
  - Per-file runtime + on_thread_start gives every alloc on a worker
    a stable, file-scoped context-id readable from sync allocator code,
    where Tokio task_locals can't reach.

Worker count matches SLT_TARGET_PARTITIONS (promoted from the previous
hardcoded 4 in test_context.rs) so a query's partition streams each
get a worker rather than contending. Total OS threads with the default
--test-threads=32 go from ~32 shared to ~32*4 + 32 = 160 per-file.
Threads are cheap (stack only); CPU work is unchanged.

The non-memory-accounting code path is preserved via cfg, so SLT runs
without the feature flag are byte-identical to today.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---
 datafusion/sqllogictest/bin/sqllogictests.rs | 36 +++++++++++++++++---
 datafusion/sqllogictest/src/accounting.rs    | 31 +++++++++++++++++
 datafusion/sqllogictest/src/lib.rs           |  7 ++--
 datafusion/sqllogictest/src/test_context.rs  |  8 ++++-
 4 files changed, 74 insertions(+), 8 deletions(-)

diff --git a/datafusion/sqllogictest/bin/sqllogictests.rs b/datafusion/sqllogictest/bin/sqllogictests.rs
index 92fb2351df957..43a4623fd8209 100644
--- a/datafusion/sqllogictest/bin/sqllogictests.rs
+++ b/datafusion/sqllogictest/bin/sqllogictests.rs
@@ -234,7 +234,7 @@ async fn run_tests() -> Result<()> {
             let currently_running_sql_tracker_clone =
                 currently_running_sql_tracker.clone();
             let file_start = Instant::now();
-            SpawnedTask::spawn(async move {
+            let body = async move {
                 let result = match (
                     options.postgres_runner,
                     options.complete,
@@ -307,9 +307,37 @@ async fn run_tests() -> Result<()> {
                 }
 
                 (result, elapsed)
-            })
-            .join()
-            .map(move |result| {
+            };
+            // Each file gets its own multi-thread runtime so a stable per-file
+            // context-id (stamped via `on_thread_start`) is readable from the
+            // global allocator hook. Bank accounting and SET-driven limit
+            // retuning will key off this id in later steps. The outer
+            // orchestration runtime hosts this via `spawn_blocking` so its
+            // worker threads aren't blocked by the per-file `block_on`.
+            //
+            // Worker count matches `SLT_TARGET_PARTITIONS` so a query's
+            // partition streams each get a worker rather than contending.
+            #[cfg(feature = "memory-accounting")]
+            let spawned = {
+                let context_id = datafusion_sqllogictest::next_context_id();
+                SpawnedTask::spawn_blocking(move || {
+                    let runtime = tokio::runtime::Builder::new_multi_thread()
+                        .enable_all()
+                        .worker_threads(datafusion_sqllogictest::SLT_TARGET_PARTITIONS)
+                        .thread_name(format!("slt-file-{context_id}"))
+                        .on_thread_start(move || {
+                            datafusion_sqllogictest::set_thread_context_id(context_id);
+                        })
+                        .build()
+                        .expect("build per-file Tokio runtime");
+                    let out = runtime.block_on(body);
+                    runtime.shutdown_background();
+                    out
+                })
+            };
+            #[cfg(not(feature = "memory-accounting"))]
+            let spawned = SpawnedTask::spawn(body);
+            spawned.join().map(move |result| {
                 let elapsed = match &result {
                     Ok((_, elapsed)) => *elapsed,
                     Err(_) => Duration::ZERO,
diff --git a/datafusion/sqllogictest/src/accounting.rs b/datafusion/sqllogictest/src/accounting.rs
index 49c4236602fba..735312eab94b4 100644
--- a/datafusion/sqllogictest/src/accounting.rs
+++ b/datafusion/sqllogictest/src/accounting.rs
@@ -69,6 +69,37 @@ static OVERDRAWN: AtomicBool = AtomicBool::new(false);
 
 thread_local! {
     static LOCAL: Cell<isize> = const { Cell::new(0) };
+
+    /// Identifies which per-file Tokio runtime owns the current worker thread.
+    /// The SLT binary stamps a fresh id onto every worker of each file's
+    /// runtime via `on_thread_start`. Zero means "not owned by a per-file
+    /// runtime" (e.g. the main thread, or the outer orchestration runtime).
+    /// Step-1 plumbing: nothing reads this yet — added so the bank lookup in
+    /// a later step has a stable per-context key.
+    static CONTEXT_ID: Cell<usize> = const { Cell::new(0) };
+}
+
+/// Monotonic source of fresh context-ids. Starts at 1; the zero value is
+/// reserved for "no per-file runtime" so callers can distinguish.
+static CONTEXT_ID_COUNTER: AtomicUsize = AtomicUsize::new(0);
+
+/// Returns a fresh, never-before-used context-id. Call once per file in the
+/// SLT binary and pass the result into the per-file runtime's
+/// `on_thread_start` callback so every worker thread of that runtime shares
+/// the same id.
+pub fn next_context_id() -> usize {
+    CONTEXT_ID_COUNTER.fetch_add(1, Ordering::Relaxed) + 1
+}
+
+/// Stamp the current thread with `id`. Intended for use inside a Tokio
+/// runtime's `on_thread_start` hook.
+pub fn set_thread_context_id(id: usize) {
+    CONTEXT_ID.with(|c| c.set(id));
+}
+
+/// Current thread's context-id, or 0 if none has been set.
+pub fn current_context_id() -> usize {
+    CONTEXT_ID.with(|c| c.get())
 }
 
 tokio::task_local! {
diff --git a/datafusion/sqllogictest/src/lib.rs b/datafusion/sqllogictest/src/lib.rs
index 5d206f5bf7b76..9248d0a8eac78 100644
--- a/datafusion/sqllogictest/src/lib.rs
+++ b/datafusion/sqllogictest/src/lib.rs
@@ -34,8 +34,9 @@ mod test_file;
 #[cfg(feature = "memory-accounting")]
 pub use accounting::{
     AccountingAllocator, OverdraftPanic, bank_balance, bytes_used, clear_overdraft,
-    init_bank, is_overdrawn, kill_on_overdraft, local_balance, memory_tracker_limit,
-    reset_bank, set_memory_tracker_limit, settle_thread_local,
+    current_context_id, init_bank, is_overdrawn, kill_on_overdraft, local_balance,
+    memory_tracker_limit, next_context_id, reset_bank, set_memory_tracker_limit,
+    set_thread_context_id, settle_thread_local,
 };
 
 pub use engines::CurrentlyExecutingSqlTracker;
@@ -56,6 +57,6 @@ mod test_context;
 mod util;
 
 pub use filters::*;
-pub use test_context::TestContext;
+pub use test_context::{SLT_TARGET_PARTITIONS, TestContext};
 pub use test_file::TestFile;
 pub use util::*;
diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs
index 590db6ddfa429..0994f0e3dcdbd 100644
--- a/datafusion/sqllogictest/src/test_context.rs
+++ b/datafusion/sqllogictest/src/test_context.rs
@@ -63,6 +63,12 @@ use log::info;
 use sqlparser::ast;
 use tempfile::TempDir;
 
+/// Target partition count used for every SLT file's `SessionConfig`. Hardcoded
+/// so query plans are deterministic across machines. The SLT binary also
+/// sizes each file's per-file Tokio runtime to this value so partition streams
+/// each get a worker rather than contending.
+pub const SLT_TARGET_PARTITIONS: usize = 4;
+
 /// Context for running tests
 pub struct TestContext {
     /// Context for running queries
@@ -123,7 +129,7 @@ impl TestContext {
     pub async fn try_new_for_test_file(relative_path: &Path) -> Option<Self> {
         let config = SessionConfig::new()
             // hardcode target partitions so plans are deterministic
-            .with_target_partitions(4);
+            .with_target_partitions(SLT_TARGET_PARTITIONS);
         let runtime = Arc::new(build_runtime_env());
 
         let mut state_builder = SessionStateBuilder::new()

From b475917ec47d9cecb4f1163987c04d5f9ce972a6 Mon Sep 17 00:00:00 2001
From: Brent Gardner <bgardner@squarelabs.net>
Date: Mon, 1 Jun 2026 11:27:38 -0600
Subject: [PATCH 11/17] WIP

---
 datafusion/sqllogictest/src/accounting.rs | 259 +++++++++++++---------
 1 file changed, 155 insertions(+), 104 deletions(-)

diff --git a/datafusion/sqllogictest/src/accounting.rs b/datafusion/sqllogictest/src/accounting.rs
index 735312eab94b4..bd6e1397da87b 100644
--- a/datafusion/sqllogictest/src/accounting.rs
+++ b/datafusion/sqllogictest/src/accounting.rs
@@ -47,35 +47,75 @@
 
 use std::alloc::{GlobalAlloc, Layout, System};
 use std::cell::Cell;
+use std::collections::HashMap;
 use std::future::Future;
 use std::sync::atomic::{AtomicBool, AtomicIsize, AtomicUsize, Ordering};
+use std::sync::{OnceLock, RwLock};
 
 /// Net byte change at which a thread flushes its local count into the bank.
 /// 64 KB chosen to keep per-thread drift tight (≤1 MB on a 16-core box) while
 /// still settling rarely enough to make the bank's atomic op amortized-free.
 const SETTLE_THRESHOLD: isize = 64 * 1024;
 
-static BANK: AtomicIsize = AtomicIsize::new(0);
+/// One account in the bank. Each SLT file's per-file Tokio runtime gets its
+/// own `Account`, keyed by the context-id stamped onto worker threads via
+/// [`set_thread_context_id`]. Context-id `0` is the catch-all for the main
+/// thread and any thread that hasn't been stamped — it never gets an account.
+struct Account {
+    /// Remaining budget — debits on alloc, credits on free, negative = overdraft.
+    balance: AtomicIsize,
+    /// Last value passed to [`init_bank`] for this account. [`bytes_used`]
+    /// expresses "allocated since budget was set" as `initial_budget - balance`,
+    /// which reads more cleanly than a deeply-negative balance.
+    initial_budget: AtomicIsize,
+    /// Sticky flag: set the moment any thread's settle observes this
+    /// account's balance crossing below zero. Cleared by [`init_bank`] or
+    /// [`clear_overdraft`].
+    overdrawn: AtomicBool,
+}
+
+impl Account {
+    fn new(budget: isize) -> Self {
+        Self {
+            balance: AtomicIsize::new(budget),
+            initial_budget: AtomicIsize::new(budget),
+            overdrawn: AtomicBool::new(false),
+        }
+    }
+}
+
+/// All known accounts, keyed by context-id. ctx-id 0 never gets an entry.
+static ACCOUNTS: OnceLock<RwLock<HashMap<usize, Account>>> = OnceLock::new();
 
-/// Last value passed to [`init_bank`]. Used by [`bytes_used`] to express
-/// "how much has been allocated since the budget was set" as a positive
-/// number, which is more readable than a deeply-negative balance.
-static INITIAL_BUDGET: AtomicIsize = AtomicIsize::new(0);
+/// Budget applied to any new account the first time its context-id is
+/// stamped onto a thread. Updated by [`init_bank`] so SLT files spawned
+/// *after* the budget is configured inherit it.
+static DEFAULT_INITIAL_BUDGET: AtomicIsize = AtomicIsize::new(0);
 
-/// Sticky flag: set the moment any thread's settle observes the bank crossing
-/// below zero. Cleared by [`init_bank`] or [`clear_overdraft`]. Reads are
-/// a single relaxed atomic load.
-static OVERDRAWN: AtomicBool = AtomicBool::new(false);
+fn accounts() -> &'static RwLock<HashMap<usize, Account>> {
+    ACCOUNTS.get_or_init(|| RwLock::new(HashMap::new()))
+}
+
+/// Run `f` against the current thread's account, or return `None` if there
+/// isn't one — silently skipping the update is fine on the alloc hot path.
+fn with_current_account<R>(f: impl FnOnce(&Account) -> R) -> Option<R> {
+    let ctx_id = CONTEXT_ID.with(|ctx| ctx.get());
+    if ctx_id == 0 {
+        return None;
+    }
+    // PERF: acquires an `RwLock` read on every settle. If it ever shows up
+    // hot, stash a `&'static Account` in a thread-local (set in
+    // `set_thread_context_id`, backed by `Box::leak`) and skip the lookup.
+    let accounts_lock = ACCOUNTS.get()?;
+    let accounts = accounts_lock.read().ok()?;
+    accounts.get(&ctx_id).map(f)
+}
 
 thread_local! {
-    static LOCAL: Cell<isize> = const { Cell::new(0) };
-
-    /// Identifies which per-file Tokio runtime owns the current worker thread.
-    /// The SLT binary stamps a fresh id onto every worker of each file's
-    /// runtime via `on_thread_start`. Zero means "not owned by a per-file
-    /// runtime" (e.g. the main thread, or the outer orchestration runtime).
-    /// Step-1 plumbing: nothing reads this yet — added so the bank lookup in
-    /// a later step has a stable per-context key.
+    static LOCAL_BALANCE: Cell<isize> = const { Cell::new(0) };
+
+    /// Account-id stamped onto worker threads via [`set_thread_context_id`].
+    /// Zero = untracked thread; nothing to track, nothing to enforce.
     static CONTEXT_ID: Cell<usize> = const { Cell::new(0) };
 }
 
@@ -91,15 +131,30 @@ pub fn next_context_id() -> usize {
     CONTEXT_ID_COUNTER.fetch_add(1, Ordering::Relaxed) + 1
 }
 
-/// Stamp the current thread with `id`. Intended for use inside a Tokio
-/// runtime's `on_thread_start` hook.
+/// Stamp the current thread with `id`. Intended for `on_thread_start`.
+/// Creates the account if it doesn't already exist.
 pub fn set_thread_context_id(id: usize) {
-    CONTEXT_ID.with(|c| c.set(id));
+    if id == 0 {
+        CONTEXT_ID.with(|ctx| ctx.set(0));
+        return;
+    }
+    // Insert under the write lock *before* stamping the thread. A HashMap
+    // resize allocates → recurses through `track` → `with_current_account`,
+    // which sees `CONTEXT_ID == 0` and bails out instead of trying to
+    // read-lock the map we're holding for writing on the same thread.
+    {
+        let accounts_lock = accounts();
+        let mut accounts = accounts_lock.write().unwrap_or_else(|p| p.into_inner());
+        accounts.entry(id).or_insert_with(|| {
+            Account::new(DEFAULT_INITIAL_BUDGET.load(Ordering::Relaxed))
+        });
+    }
+    CONTEXT_ID.with(|ctx| ctx.set(id));
 }
 
 /// Current thread's context-id, or 0 if none has been set.
 pub fn current_context_id() -> usize {
-    CONTEXT_ID.with(|c| c.get())
+    CONTEXT_ID.with(|ctx| ctx.get())
 }
 
 tokio::task_local! {
@@ -137,36 +192,34 @@ pub async fn kill_on_overdraft<F: Future>(f: F) -> F::Output {
     KILL_GUARD.scope(Cell::new(false), f).await
 }
 
-/// Set the bank to `value`, remember it as the budget for [`bytes_used`],
-/// and clear the overdraft flag. Idempotent; safe to call from any thread.
-///
-/// Typical use: call once early in `main` to preload a memory budget before
-/// significant allocation occurs. Note that per-thread locals are NOT zeroed,
-/// so calling this mid-program leaves prior unsettled deltas in the locals
-/// that will eventually flush into the bank.
+/// Set the default budget for accounts created later, and reset the
+/// current thread's account (if any) to `value` with cleared overdraft.
 pub fn init_bank(value: isize) {
-    INITIAL_BUDGET.store(value, Ordering::Relaxed);
-    BANK.store(value, Ordering::Relaxed);
-    OVERDRAWN.store(false, Ordering::Relaxed);
+    DEFAULT_INITIAL_BUDGET.store(value, Ordering::Relaxed);
+    let _ = with_current_account(|acct| {
+        acct.initial_budget.store(value, Ordering::Relaxed);
+        acct.balance.store(value, Ordering::Relaxed);
+        acct.overdrawn.store(false, Ordering::Relaxed);
+    });
 }
 
-/// Bytes allocated since the last [`init_bank`] — i.e., `budget - bank_balance`.
-/// Always positive (or zero) under normal conditions; only meaningful relative
-/// to the most recent `init_bank` call.
+/// Bytes allocated against the *current thread's* context since its budget
+/// was last set — i.e., `initial_budget - bank`. Always positive (or zero)
+/// under normal conditions; returns `0` if the bank map isn't initialized.
 pub fn bytes_used() -> isize {
-    INITIAL_BUDGET.load(Ordering::Relaxed) - BANK.load(Ordering::Relaxed)
+    with_current_account(|acct| {
+        acct.initial_budget.load(Ordering::Relaxed) - acct.balance.load(Ordering::Relaxed)
+    })
+    .unwrap_or(0)
 }
 
-/// Reset the bank to the budget last passed to [`init_bank`] and clear the
-/// overdraft flag. No-op if `init_bank` has never been called (budget = 0).
-///
-/// Intended to be called between independent units of work — e.g., between
-/// SQL statements — so each gets a fresh budget regardless of whether the
-/// previous one drifted the bank or tripped overdraft.
+/// Reset the current thread's account to its initial budget; clear overdraft.
 pub fn reset_bank() {
-    let budget = INITIAL_BUDGET.load(Ordering::Relaxed);
-    BANK.store(budget, Ordering::Relaxed);
-    OVERDRAWN.store(false, Ordering::Relaxed);
+    let _ = with_current_account(|acct| {
+        let budget = acct.initial_budget.load(Ordering::Relaxed);
+        acct.balance.store(budget, Ordering::Relaxed);
+        acct.overdrawn.store(false, Ordering::Relaxed);
+    });
 }
 
 /// Cross-module config for DataFusion's voluntary `MemoryPool` limit, set
@@ -185,61 +238,71 @@ pub fn memory_tracker_limit() -> usize {
     MEMORY_TRACKER_LIMIT.load(Ordering::Relaxed)
 }
 
-/// True iff any settle has ever seen the bank go below zero since the last
-/// [`init_bank`] / [`clear_overdraft`]. Single relaxed atomic load.
+/// True iff any settle has ever seen the *current thread's* context bank
+/// go below zero since the last [`init_bank`] / [`clear_overdraft`]. Returns
+/// `false` if the bank map isn't initialized or the current context has no
+/// entry.
 pub fn is_overdrawn() -> bool {
-    OVERDRAWN.load(Ordering::Relaxed)
+    with_current_account(|acct| acct.overdrawn.load(Ordering::Relaxed)).unwrap_or(false)
 }
 
-/// Clear the sticky overdraft flag without touching the bank.
+/// Clear the *current thread's* context overdraft flag without touching
+/// its bank.
 pub fn clear_overdraft() {
-    OVERDRAWN.store(false, Ordering::Relaxed);
+    let _ = with_current_account(|acct| acct.overdrawn.store(false, Ordering::Relaxed));
 }
 
-/// Current bank balance. Allocs debit, deallocs credit, negative = overdraft.
-/// Single relaxed atomic load; cheap to call from any thread.
-///
-/// Lags reality by up to `SETTLE_THRESHOLD * num_threads`. To get an exact
-/// snapshot, drain each thread's local first via [`settle_thread_local`] —
-/// usually overkill.
+/// Current account balance. Negative = overdraft. `0` if untracked.
 pub fn bank_balance() -> isize {
-    BANK.load(Ordering::Relaxed)
+    with_current_account(|acct| acct.balance.load(Ordering::Relaxed)).unwrap_or(0)
 }
 
 /// Current thread's local balance — not yet reflected in the global bank.
 /// Always in `(-SETTLE_THRESHOLD, +SETTLE_THRESHOLD)`. Sign matches the bank:
 /// negative on a thread that's net-allocated, positive on one that's net-freed.
 pub fn local_balance() -> isize {
-    LOCAL.with(|c| c.get())
+    LOCAL_BALANCE.with(|bal| bal.get())
 }
 
-/// Force the current thread to flush its local count into the bank.
+/// Force the current thread to flush its local count into its context bank.
+/// No-op on untracked threads (`CONTEXT_ID == 0`).
 pub fn settle_thread_local() {
-    let _ = LOCAL.try_with(|c| {
-        let v = c.replace(0);
+    if CONTEXT_ID.with(|ctx| ctx.get()) == 0 {
+        return;
+    }
+    let _ = LOCAL_BALANCE.try_with(|bal| {
+        let v = bal.replace(0);
         if v != 0 {
-            BANK.fetch_add(v, Ordering::Relaxed);
+            let _ = with_current_account(|acct| {
+                acct.balance.fetch_add(v, Ordering::Relaxed);
+            });
         }
     });
 }
 
-/// Record a delta against the current thread's local; flush to bank if we've
-/// crossed `±SETTLE_THRESHOLD`; on a debit, possibly kill the task. Inlines
-/// into the alloc hot path.
+/// Record a delta into the current thread's account; settle when local drift
+/// crosses `±SETTLE_THRESHOLD`; kill the task if a debit hits overdraft.
+/// Inlines into the alloc hot path.
 #[inline(always)]
 fn track(delta: isize) {
-    let _ = LOCAL.try_with(|c| {
-        let new = c.get() + delta;
+    // Untracked thread — skip everything.
+    if CONTEXT_ID.with(|ctx| ctx.get()) == 0 {
+        return;
+    }
+    let _ = LOCAL_BALANCE.try_with(|bal| {
+        let new = bal.get() + delta;
         if new >= SETTLE_THRESHOLD || new <= -SETTLE_THRESHOLD {
-            // `fetch_add` returns the OLD value. new_bank = old + new.
-            let old_bank = BANK.fetch_add(new, Ordering::Relaxed);
-            let new_bank = old_bank.wrapping_add(new);
-            if new_bank < 0 {
-                OVERDRAWN.store(true, Ordering::Relaxed);
-            }
-            c.set(0);
+            let _ = with_current_account(|acct| {
+                // `fetch_add` returns the OLD value. new_balance = old + new.
+                let old_balance = acct.balance.fetch_add(new, Ordering::Relaxed);
+                let new_balance = old_balance.wrapping_add(new);
+                if new_balance < 0 {
+                    acct.overdrawn.store(true, Ordering::Relaxed);
+                }
+            });
+            bal.set(0);
         } else {
-            c.set(new);
+            bal.set(new);
         }
     });
 
@@ -247,7 +310,9 @@ fn track(delta: isize) {
     // code paths (alloc / alloc_zeroed / growing realloc). Credits happen
     // during dealloc which runs inside Drop during unwinding — a panic there
     // would double-fault and abort the process.
-    if delta < 0 && OVERDRAWN.load(Ordering::Relaxed) {
+    if delta < 0
+        && with_current_account(|acct| acct.overdrawn.load(Ordering::Relaxed)).unwrap_or(false)
+    {
         maybe_kill();
     }
 }
@@ -260,7 +325,8 @@ fn maybe_kill() {
     let _ = KILL_GUARD.try_with(|panicking| {
         if !panicking.get() {
             panicking.set(true);
-            let balance = BANK.load(Ordering::Relaxed);
+            let balance =
+                with_current_account(|acct| acct.balance.load(Ordering::Relaxed)).unwrap_or(0);
             std::panic::panic_any(OverdraftPanic {
                 bank_balance: balance,
             });
@@ -331,22 +397,20 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> {
 mod tests {
     use super::*;
     use futures::FutureExt;
-    use std::sync::{Mutex, MutexGuard};
 
     #[global_allocator]
     static GLOBAL: AccountingAllocator = AccountingAllocator::system();
 
-    /// All accounting tests share the global bank, so they must run serially.
-    /// Each test takes this lock as its first line — a panicked test poisons
-    /// the lock; we recover the guard so subsequent tests still run.
-    static SERIAL: Mutex<()> = Mutex::new(());
-    fn serial() -> MutexGuard<'static, ()> {
-        SERIAL.lock().unwrap_or_else(|p| p.into_inner())
+    /// Each test runs on its own thread (cargo-test parallelism) and stamps a
+    /// fresh context-id, so per-context isolation makes them naturally
+    /// independent — no shared mutex required.
+    fn enter_fresh_context() {
+        set_thread_context_id(next_context_id());
     }
 
     #[test]
     fn alloc_debits_and_free_credits_bank() {
-        let _g = serial();
+        enter_fresh_context();
         settle_thread_local();
         let before = bank_balance();
 
@@ -371,21 +435,20 @@ mod tests {
 
     #[test]
     fn init_bank_sets_value() {
-        let _g = serial();
+        enter_fresh_context();
         init_bank(1_000_000);
-        // The bank drifts from concurrent allocator activity in other tests
-        // (allocs debit), so we expect at-or-below the set value.
+        // The bank drifts a little from this thread's own allocator activity
+        // between init and the read, so we expect at-or-below the set value.
         let v = bank_balance();
         assert!(
             (900_000..=1_000_000).contains(&v),
             "init_bank didn't stick: v={v}"
         );
-        init_bank(0);
     }
 
     #[test]
     fn overdraft_flag_trips_when_bank_goes_negative() {
-        let _g = serial();
+        enter_fresh_context();
         init_bank(1024);
         assert!(!is_overdrawn(), "flag should be clear after init_bank");
 
@@ -405,20 +468,11 @@ mod tests {
 
         clear_overdraft();
         assert!(!is_overdrawn(), "clear_overdraft didn't clear the flag");
-
-        init_bank(0);
     }
 
-    // Holding the std Mutex guard across `.await` is fine here: the test uses
-    // a current-thread runtime so there's no other task to block, and we need
-    // exclusive access to the global bank for the entire test body.
-    #[expect(
-        clippy::await_holding_lock,
-        reason = "test uses current-thread runtime; the guard exists to serialize tests against the shared global bank"
-    )]
     #[tokio::test(flavor = "current_thread")]
     async fn kill_on_overdraft_panics_inside_scope() {
-        let _g = serial();
+        enter_fresh_context();
         init_bank(1024);
 
         let result = std::panic::AssertUnwindSafe(kill_on_overdraft(async {
@@ -442,13 +496,11 @@ mod tests {
             "payload should report negative balance; got {}",
             overdraft.bank_balance
         );
-
-        init_bank(0);
     }
 
     #[tokio::test(flavor = "current_thread")]
     async fn no_kill_outside_scope() {
-        let _g = serial();
+        enter_fresh_context();
         // No `kill_on_overdraft` wrapper -- overdraft should set the flag
         // but allocations keep proceeding without panicking.
         init_bank(1024);
@@ -456,12 +508,11 @@ mod tests {
         let _v2: Vec<u8> = vec![0u8; 1024];
         assert!(is_overdrawn(), "should have overdrawn");
         // Reaching here without panic is the whole point.
-        init_bank(0);
     }
 
     #[test]
     fn threshold_settlement_flushes_to_bank() {
-        let _g = serial();
+        enter_fresh_context();
         settle_thread_local();
         let before_bank = bank_balance();
 

From 79130259e77f91531dc64286597ee8293afaddf0 Mon Sep 17 00:00:00 2001
From: Brent Gardner <bgardner@squarelabs.net>
Date: Mon, 1 Jun 2026 11:48:51 -0600
Subject: [PATCH 12/17] WIP2

---
 datafusion/sqllogictest/src/accounting.rs     | 191 ++++++------------
 .../src/engines/datafusion_engine/runner.rs   |  26 +--
 datafusion/sqllogictest/src/lib.rs            |   7 +-
 3 files changed, 71 insertions(+), 153 deletions(-)

diff --git a/datafusion/sqllogictest/src/accounting.rs b/datafusion/sqllogictest/src/accounting.rs
index bd6e1397da87b..d2741db901c3f 100644
--- a/datafusion/sqllogictest/src/accounting.rs
+++ b/datafusion/sqllogictest/src/accounting.rs
@@ -49,7 +49,7 @@ use std::alloc::{GlobalAlloc, Layout, System};
 use std::cell::Cell;
 use std::collections::HashMap;
 use std::future::Future;
-use std::sync::atomic::{AtomicBool, AtomicIsize, AtomicUsize, Ordering};
+use std::sync::atomic::{AtomicIsize, AtomicUsize, Ordering};
 use std::sync::{OnceLock, RwLock};
 
 /// Net byte change at which a thread flushes its local count into the bank.
@@ -57,54 +57,29 @@ use std::sync::{OnceLock, RwLock};
 /// still settling rarely enough to make the bank's atomic op amortized-free.
 const SETTLE_THRESHOLD: isize = 64 * 1024;
 
-/// One account in the bank. Each SLT file's per-file Tokio runtime gets its
-/// own `Account`, keyed by the context-id stamped onto worker threads via
-/// [`set_thread_context_id`]. Context-id `0` is the catch-all for the main
-/// thread and any thread that hasn't been stamped — it never gets an account.
-struct Account {
-    /// Remaining budget — debits on alloc, credits on free, negative = overdraft.
-    balance: AtomicIsize,
-    /// Last value passed to [`init_bank`] for this account. [`bytes_used`]
-    /// expresses "allocated since budget was set" as `initial_budget - balance`,
-    /// which reads more cleanly than a deeply-negative balance.
-    initial_budget: AtomicIsize,
-    /// Sticky flag: set the moment any thread's settle observes this
-    /// account's balance crossing below zero. Cleared by [`init_bank`] or
-    /// [`clear_overdraft`].
-    overdrawn: AtomicBool,
-}
-
-impl Account {
-    fn new(budget: isize) -> Self {
-        Self {
-            balance: AtomicIsize::new(budget),
-            initial_budget: AtomicIsize::new(budget),
-            overdrawn: AtomicBool::new(false),
-        }
-    }
-}
-
-/// All known accounts, keyed by context-id. ctx-id 0 never gets an entry.
-static ACCOUNTS: OnceLock<RwLock<HashMap<usize, Account>>> = OnceLock::new();
+/// The bank: every account, keyed by context-id, valued by remaining budget.
+/// Debits on alloc, credits on free, negative = overdraft. ctx-id 0 never
+/// gets an entry — that's the "untracked thread" marker.
+static ACCOUNTS: OnceLock<RwLock<HashMap<usize, AtomicIsize>>> = OnceLock::new();
 
-/// Budget applied to any new account the first time its context-id is
-/// stamped onto a thread. Updated by [`init_bank`] so SLT files spawned
-/// *after* the budget is configured inherit it.
-static DEFAULT_INITIAL_BUDGET: AtomicIsize = AtomicIsize::new(0);
+/// Starting budget for any new account, set by [`init_bank`] and inherited
+/// by per-file SLT contexts spawned after.
+static DEFAULT_BUDGET: AtomicIsize = AtomicIsize::new(0);
 
-fn accounts() -> &'static RwLock<HashMap<usize, Account>> {
+fn accounts() -> &'static RwLock<HashMap<usize, AtomicIsize>> {
     ACCOUNTS.get_or_init(|| RwLock::new(HashMap::new()))
 }
 
-/// Run `f` against the current thread's account, or return `None` if there
-/// isn't one — silently skipping the update is fine on the alloc hot path.
-fn with_current_account<R>(f: impl FnOnce(&Account) -> R) -> Option<R> {
+/// Run `f` against the current thread's account balance, or return `None`
+/// if there isn't one — silently skipping the update is fine on the alloc
+/// hot path.
+fn with_current_balance<R>(f: impl FnOnce(&AtomicIsize) -> R) -> Option<R> {
     let ctx_id = CONTEXT_ID.with(|ctx| ctx.get());
     if ctx_id == 0 {
         return None;
     }
     // PERF: acquires an `RwLock` read on every settle. If it ever shows up
-    // hot, stash a `&'static Account` in a thread-local (set in
+    // hot, stash a `&'static AtomicIsize` in a thread-local (set in
     // `set_thread_context_id`, backed by `Box::leak`) and skip the lookup.
     let accounts_lock = ACCOUNTS.get()?;
     let accounts = accounts_lock.read().ok()?;
@@ -145,9 +120,9 @@ pub fn set_thread_context_id(id: usize) {
     {
         let accounts_lock = accounts();
         let mut accounts = accounts_lock.write().unwrap_or_else(|p| p.into_inner());
-        accounts.entry(id).or_insert_with(|| {
-            Account::new(DEFAULT_INITIAL_BUDGET.load(Ordering::Relaxed))
-        });
+        accounts
+            .entry(id)
+            .or_insert_with(|| AtomicIsize::new(DEFAULT_BUDGET.load(Ordering::Relaxed)));
     }
     CONTEXT_ID.with(|ctx| ctx.set(id));
 }
@@ -193,33 +168,10 @@ pub async fn kill_on_overdraft<F: Future>(f: F) -> F::Output {
 }
 
 /// Set the default budget for accounts created later, and reset the
-/// current thread's account (if any) to `value` with cleared overdraft.
+/// current thread's account (if any) to `value`.
 pub fn init_bank(value: isize) {
-    DEFAULT_INITIAL_BUDGET.store(value, Ordering::Relaxed);
-    let _ = with_current_account(|acct| {
-        acct.initial_budget.store(value, Ordering::Relaxed);
-        acct.balance.store(value, Ordering::Relaxed);
-        acct.overdrawn.store(false, Ordering::Relaxed);
-    });
-}
-
-/// Bytes allocated against the *current thread's* context since its budget
-/// was last set — i.e., `initial_budget - bank`. Always positive (or zero)
-/// under normal conditions; returns `0` if the bank map isn't initialized.
-pub fn bytes_used() -> isize {
-    with_current_account(|acct| {
-        acct.initial_budget.load(Ordering::Relaxed) - acct.balance.load(Ordering::Relaxed)
-    })
-    .unwrap_or(0)
-}
-
-/// Reset the current thread's account to its initial budget; clear overdraft.
-pub fn reset_bank() {
-    let _ = with_current_account(|acct| {
-        let budget = acct.initial_budget.load(Ordering::Relaxed);
-        acct.balance.store(budget, Ordering::Relaxed);
-        acct.overdrawn.store(false, Ordering::Relaxed);
-    });
+    DEFAULT_BUDGET.store(value, Ordering::Relaxed);
+    let _ = with_current_balance(|bal| bal.store(value, Ordering::Relaxed));
 }
 
 /// Cross-module config for DataFusion's voluntary `MemoryPool` limit, set
@@ -238,23 +190,9 @@ pub fn memory_tracker_limit() -> usize {
     MEMORY_TRACKER_LIMIT.load(Ordering::Relaxed)
 }
 
-/// True iff any settle has ever seen the *current thread's* context bank
-/// go below zero since the last [`init_bank`] / [`clear_overdraft`]. Returns
-/// `false` if the bank map isn't initialized or the current context has no
-/// entry.
-pub fn is_overdrawn() -> bool {
-    with_current_account(|acct| acct.overdrawn.load(Ordering::Relaxed)).unwrap_or(false)
-}
-
-/// Clear the *current thread's* context overdraft flag without touching
-/// its bank.
-pub fn clear_overdraft() {
-    let _ = with_current_account(|acct| acct.overdrawn.store(false, Ordering::Relaxed));
-}
-
 /// Current account balance. Negative = overdraft. `0` if untracked.
 pub fn bank_balance() -> isize {
-    with_current_account(|acct| acct.balance.load(Ordering::Relaxed)).unwrap_or(0)
+    with_current_balance(|bal| bal.load(Ordering::Relaxed)).unwrap_or(0)
 }
 
 /// Current thread's local balance — not yet reflected in the global bank.
@@ -273,16 +211,19 @@ pub fn settle_thread_local() {
     let _ = LOCAL_BALANCE.try_with(|bal| {
         let v = bal.replace(0);
         if v != 0 {
-            let _ = with_current_account(|acct| {
-                acct.balance.fetch_add(v, Ordering::Relaxed);
-            });
+            let _ = with_current_balance(|b| b.fetch_add(v, Ordering::Relaxed));
         }
     });
 }
 
 /// Record a delta into the current thread's account; settle when local drift
-/// crosses `±SETTLE_THRESHOLD`; kill the task if a debit hits overdraft.
-/// Inlines into the alloc hot path.
+/// crosses `±SETTLE_THRESHOLD`; kill the task if a debit settle leaves the
+/// balance negative. Inlines into the alloc hot path.
+///
+/// Only debits can fire the kill, and only debits run inside Drop-free
+/// code paths (alloc / alloc_zeroed / growing realloc). Credits happen
+/// during dealloc which runs inside Drop during unwinding — a panic there
+/// would double-fault and abort the process.
 #[inline(always)]
 fn track(delta: isize) {
     // Untracked thread — skip everything.
@@ -292,44 +233,38 @@ fn track(delta: isize) {
     let _ = LOCAL_BALANCE.try_with(|bal| {
         let new = bal.get() + delta;
         if new >= SETTLE_THRESHOLD || new <= -SETTLE_THRESHOLD {
-            let _ = with_current_account(|acct| {
-                // `fetch_add` returns the OLD value. new_balance = old + new.
-                let old_balance = acct.balance.fetch_add(new, Ordering::Relaxed);
-                let new_balance = old_balance.wrapping_add(new);
-                if new_balance < 0 {
-                    acct.overdrawn.store(true, Ordering::Relaxed);
-                }
+            // Settle and drop the read lock *before* calling `maybe_kill`.
+            // The kill panic allocates its payload, which recurses through
+            // `track` and would self-deadlock on `std::sync::RwLock` if we
+            // were still holding it on this thread.
+            let new_balance = with_current_balance(|b| {
+                // `fetch_add` returns the OLD value.
+                b.fetch_add(new, Ordering::Relaxed).wrapping_add(new)
             });
             bal.set(0);
+            if new < 0 {
+                if let Some(nb) = new_balance {
+                    if nb < 0 {
+                        maybe_kill(nb);
+                    }
+                }
+            }
         } else {
             bal.set(new);
         }
     });
-
-    // Only debits can drive overdraft, and only debits run inside Drop-free
-    // code paths (alloc / alloc_zeroed / growing realloc). Credits happen
-    // during dealloc which runs inside Drop during unwinding — a panic there
-    // would double-fault and abort the process.
-    if delta < 0
-        && with_current_account(|acct| acct.overdrawn.load(Ordering::Relaxed)).unwrap_or(false)
-    {
-        maybe_kill();
-    }
 }
 
 /// If the current task is inside a [`kill_on_overdraft`] scope and we haven't
-/// already started panicking for it, fire the kill panic.
+/// already started panicking for it, fire the kill panic with `bank_balance`
+/// as the payload.
 #[cold]
 #[inline(never)]
-fn maybe_kill() {
+fn maybe_kill(bank_balance: isize) {
     let _ = KILL_GUARD.try_with(|panicking| {
         if !panicking.get() {
             panicking.set(true);
-            let balance =
-                with_current_account(|acct| acct.balance.load(Ordering::Relaxed)).unwrap_or(0);
-            std::panic::panic_any(OverdraftPanic {
-                bank_balance: balance,
-            });
+            std::panic::panic_any(OverdraftPanic { bank_balance });
         }
     });
 }
@@ -447,27 +382,19 @@ mod tests {
     }
 
     #[test]
-    fn overdraft_flag_trips_when_bank_goes_negative() {
+    fn bank_goes_negative_on_overdraft() {
         enter_fresh_context();
         init_bank(1024);
-        assert!(!is_overdrawn(), "flag should be clear after init_bank");
+        assert!(bank_balance() > 0, "should start positive after init_bank");
 
         // Alloc large enough to cross SETTLE_THRESHOLD in one shot — settle
-        // drives bank from +1024 to roughly -(threshold+budget+overhead).
+        // drives balance from +1024 to roughly -(threshold+budget+overhead).
         let _v: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 4096];
-        assert!(
-            is_overdrawn(),
-            "overdraft flag didn't trip; bank={}",
-            bank_balance()
-        );
         assert!(
             bank_balance() < 0,
             "bank should be negative after overdraft; bank={}",
             bank_balance()
         );
-
-        clear_overdraft();
-        assert!(!is_overdrawn(), "clear_overdraft didn't clear the flag");
     }
 
     #[tokio::test(flavor = "current_thread")]
@@ -476,12 +403,9 @@ mod tests {
         init_bank(1024);
 
         let result = std::panic::AssertUnwindSafe(kill_on_overdraft(async {
-            // First debit larger than SETTLE_THRESHOLD auto-settles and
-            // drives the bank to negative → next debit-with-overdraft kills.
-            // We do two big allocations so the second is the one that
-            // observes OVERDRAWN=true and panics.
-            let _v1: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 4096];
-            let _v2: Vec<u8> = vec![0u8; 1024];
+            // Big alloc crosses SETTLE_THRESHOLD → settle drives balance
+            // negative → maybe_kill fires inside the settle closure → panic.
+            let _v: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 4096];
             unreachable!("should have panicked");
         }))
         .catch_unwind()
@@ -501,12 +425,11 @@ mod tests {
     #[tokio::test(flavor = "current_thread")]
     async fn no_kill_outside_scope() {
         enter_fresh_context();
-        // No `kill_on_overdraft` wrapper -- overdraft should set the flag
-        // but allocations keep proceeding without panicking.
+        // No `kill_on_overdraft` wrapper -- overdraft should drive the bank
+        // negative but allocations keep proceeding without panicking.
         init_bank(1024);
-        let _v1: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 4096];
-        let _v2: Vec<u8> = vec![0u8; 1024];
-        assert!(is_overdrawn(), "should have overdrawn");
+        let _v: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 4096];
+        assert!(bank_balance() < 0, "should have overdrawn");
         // Reaching here without panic is the whole point.
     }
 
diff --git a/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs b/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs
index 5e35b9c981e40..82889d17c911a 100644
--- a/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs
+++ b/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs
@@ -90,14 +90,9 @@ impl DataFusion {
     async fn run_one(&self, sql: &str) -> Result<DFOutput> {
         #[cfg(feature = "memory-accounting")]
         {
-            use crate::{OverdraftPanic, kill_on_overdraft, reset_bank};
+            use crate::{OverdraftPanic, kill_on_overdraft};
             use futures::FutureExt;
 
-            // Fresh budget for this statement. Independent of whatever the
-            // previous statement did (or panicked through). Also clears the
-            // sticky overdraft flag.
-            reset_bank();
-
             let fut = kill_on_overdraft(run_query(
                 &self.ctx,
                 is_spark_path(&self.relative_path),
@@ -108,22 +103,23 @@ impl DataFusion {
                 Ok(r) => r,
                 Err(payload) => {
                     if let Some(od) = payload.downcast_ref::<OverdraftPanic>() {
-                        let actual = (crate::bytes_used().max(0) as u64) / (1024 * 1024);
-                        let supposed = (self.ctx.runtime_env().memory_pool.reserved()
+                        let df_reserved_mb = (self
+                            .ctx
+                            .runtime_env()
+                            .memory_pool
+                            .reserved()
                             as u64)
                             / (1024 * 1024);
                         warn!(
                             "[{}] killed by allocator overdraft: \
-                             actual={actual} MB, supposed={supposed} MB \
-                             (bank balance at panic = {} bytes); sql={sql:?}",
+                             bank balance = {} bytes, df-pool reserved = {df_reserved_mb} MB; \
+                             sql = {sql:?}",
                             self.relative_path.display(),
-                            od.bank_balance
+                            od.bank_balance,
                         );
-                        // Pre-emptive reset so the next statement starts clean
-                        // even if some drop chain didn't fully credit the bank.
-                        reset_bank();
                         Err(DFSqlLogicTestError::Other(format!(
-                            "allocator overdraft: actual={actual} MB supposed={supposed} MB"
+                            "allocator overdraft: bank balance at panic = {} bytes",
+                            od.bank_balance,
                         )))
                     } else {
                         // Not our panic — re-raise so test runner sees it.
diff --git a/datafusion/sqllogictest/src/lib.rs b/datafusion/sqllogictest/src/lib.rs
index 9248d0a8eac78..00735dd92f4b8 100644
--- a/datafusion/sqllogictest/src/lib.rs
+++ b/datafusion/sqllogictest/src/lib.rs
@@ -33,10 +33,9 @@ mod test_file;
 
 #[cfg(feature = "memory-accounting")]
 pub use accounting::{
-    AccountingAllocator, OverdraftPanic, bank_balance, bytes_used, clear_overdraft,
-    current_context_id, init_bank, is_overdrawn, kill_on_overdraft, local_balance,
-    memory_tracker_limit, next_context_id, reset_bank, set_memory_tracker_limit,
-    set_thread_context_id, settle_thread_local,
+    AccountingAllocator, OverdraftPanic, bank_balance, current_context_id, init_bank,
+    kill_on_overdraft, local_balance, memory_tracker_limit, next_context_id,
+    set_memory_tracker_limit, set_thread_context_id, settle_thread_local,
 };
 
 pub use engines::CurrentlyExecutingSqlTracker;

From 815e7461992e5ad1ab758beabe7f14fc63343005 Mon Sep 17 00:00:00 2001
From: Brent Gardner <bgardner@squarelabs.net>
Date: Mon, 1 Jun 2026 12:22:17 -0600
Subject: [PATCH 13/17] WIP3

---
 datafusion/sqllogictest/src/accounting.rs | 136 +++++++++++-----------
 1 file changed, 65 insertions(+), 71 deletions(-)

diff --git a/datafusion/sqllogictest/src/accounting.rs b/datafusion/sqllogictest/src/accounting.rs
index d2741db901c3f..78aacec0c228d 100644
--- a/datafusion/sqllogictest/src/accounting.rs
+++ b/datafusion/sqllogictest/src/accounting.rs
@@ -73,7 +73,7 @@ fn accounts() -> &'static RwLock<HashMap<usize, AtomicIsize>> {
 /// Run `f` against the current thread's account balance, or return `None`
 /// if there isn't one — silently skipping the update is fine on the alloc
 /// hot path.
-fn with_current_balance<R>(f: impl FnOnce(&AtomicIsize) -> R) -> Option<R> {
+fn with_current_balance<R>(op: impl FnOnce(&AtomicIsize) -> R) -> Option<R> {
     let ctx_id = CONTEXT_ID.with(|ctx| ctx.get());
     if ctx_id == 0 {
         return None;
@@ -83,7 +83,7 @@ fn with_current_balance<R>(f: impl FnOnce(&AtomicIsize) -> R) -> Option<R> {
     // `set_thread_context_id`, backed by `Box::leak`) and skip the lookup.
     let accounts_lock = ACCOUNTS.get()?;
     let accounts = accounts_lock.read().ok()?;
-    accounts.get(&ctx_id).map(f)
+    accounts.get(&ctx_id).map(op)
 }
 
 thread_local! {
@@ -119,7 +119,7 @@ pub fn set_thread_context_id(id: usize) {
     // read-lock the map we're holding for writing on the same thread.
     {
         let accounts_lock = accounts();
-        let mut accounts = accounts_lock.write().unwrap_or_else(|p| p.into_inner());
+        let mut accounts = accounts_lock.write().unwrap_or_else(|poison| poison.into_inner());
         accounts
             .entry(id)
             .or_insert_with(|| AtomicIsize::new(DEFAULT_BUDGET.load(Ordering::Relaxed)));
@@ -133,12 +133,11 @@ pub fn current_context_id() -> usize {
 }
 
 tokio::task_local! {
-    /// Active iff the current task is inside a [`kill_on_overdraft`] scope.
-    /// Holding a `Cell<bool>` lets us track "we already started panicking"
-    /// without a separate thread-local, since the cell is task-scoped — it
-    /// disappears when the future unwinds past the scope, so the next task
-    /// on the same worker thread sees no enforcement.
-    static KILL_GUARD: Cell<bool>;
+    /// Set iff the current task is inside a [`kill_on_overdraft`] scope, where
+    /// overdrafts are disallowed and the allocator must panic the task on
+    /// the first one. The `Cell<bool>` tracks whether the panic has already
+    /// fired so drop-chain allocations during unwinding don't double-fault.
+    static DISALLOW_OVERDRAFT: Cell<bool>;
 }
 
 /// Payload attached to allocator-induced panics. Catch with:
@@ -163,8 +162,8 @@ pub struct OverdraftPanic {
 /// unwinding past the scope, the worker thread reverts to the default
 /// "count, but don't kill" mode, so any non-marked futures it polls next
 /// (system tasks, hyper handlers, etc.) are unaffected.
-pub async fn kill_on_overdraft<F: Future>(f: F) -> F::Output {
-    KILL_GUARD.scope(Cell::new(false), f).await
+pub async fn kill_on_overdraft<F: Future>(future: F) -> F::Output {
+    DISALLOW_OVERDRAFT.scope(Cell::new(false), future).await
 }
 
 /// Set the default budget for accounts created later, and reset the
@@ -199,7 +198,7 @@ pub fn bank_balance() -> isize {
 /// Always in `(-SETTLE_THRESHOLD, +SETTLE_THRESHOLD)`. Sign matches the bank:
 /// negative on a thread that's net-allocated, positive on one that's net-freed.
 pub fn local_balance() -> isize {
-    LOCAL_BALANCE.with(|bal| bal.get())
+    LOCAL_BALANCE.with(|loc_bal| loc_bal.get())
 }
 
 /// Force the current thread to flush its local count into its context bank.
@@ -208,62 +207,57 @@ pub fn settle_thread_local() {
     if CONTEXT_ID.with(|ctx| ctx.get()) == 0 {
         return;
     }
-    let _ = LOCAL_BALANCE.try_with(|bal| {
-        let v = bal.replace(0);
-        if v != 0 {
-            let _ = with_current_balance(|b| b.fetch_add(v, Ordering::Relaxed));
+    let _ = LOCAL_BALANCE.try_with(|loc_bal| {
+        let drift = loc_bal.replace(0);
+        if drift != 0 {
+            let _ = with_current_balance(|bal| bal.fetch_add(drift, Ordering::Relaxed));
         }
     });
 }
 
-/// Record a delta into the current thread's account; settle when local drift
-/// crosses `±SETTLE_THRESHOLD`; kill the task if a debit settle leaves the
-/// balance negative. Inlines into the alloc hot path.
-///
-/// Only debits can fire the kill, and only debits run inside Drop-free
-/// code paths (alloc / alloc_zeroed / growing realloc). Credits happen
-/// during dealloc which runs inside Drop during unwinding — a panic there
-/// would double-fault and abort the process.
+/// Record a delta into the current thread's account: settle local drift into
+/// the bank when it crosses `±SETTLE_THRESHOLD`, fire the kill panic on a
+/// debit that leaves the account negative.
 #[inline(always)]
 fn track(delta: isize) {
-    // Untracked thread — skip everything.
     if CONTEXT_ID.with(|ctx| ctx.get()) == 0 {
         return;
     }
-    let _ = LOCAL_BALANCE.try_with(|bal| {
-        let new = bal.get() + delta;
-        if new >= SETTLE_THRESHOLD || new <= -SETTLE_THRESHOLD {
-            // Settle and drop the read lock *before* calling `maybe_kill`.
-            // The kill panic allocates its payload, which recurses through
-            // `track` and would self-deadlock on `std::sync::RwLock` if we
-            // were still holding it on this thread.
-            let new_balance = with_current_balance(|b| {
-                // `fetch_add` returns the OLD value.
-                b.fetch_add(new, Ordering::Relaxed).wrapping_add(new)
-            });
-            bal.set(0);
-            if new < 0 {
-                if let Some(nb) = new_balance {
-                    if nb < 0 {
-                        maybe_kill(nb);
-                    }
-                }
-            }
-        } else {
-            bal.set(new);
+    let _ = LOCAL_BALANCE.try_with(|loc_bal| {
+        let drift = loc_bal.get() + delta;
+        // 99% case: drift fits — accumulate locally and bail.
+        if -SETTLE_THRESHOLD < drift && drift < SETTLE_THRESHOLD {
+            loc_bal.set(drift);
+            return;
+        }
+        // Drop the read lock *before* maybe_kill — the panic allocates,
+        // recurses through track, and would self-deadlock on std::sync::RwLock.
+        let new_bal = with_current_balance(|bal| {
+            bal.fetch_add(drift, Ordering::Relaxed).wrapping_add(drift)
+        });
+        loc_bal.set(0);
+        // Only debits fire the kill — credits run inside Drop chains during
+        // unwinding, where a panic would double-fault and abort the process.
+        if delta >= 0 {
+            return;
+        }
+        let Some(new_bal) = new_bal else { return };
+        if new_bal >= 0 {
+            return;
         }
+        maybe_kill(new_bal);
     });
 }
 
-/// If the current task is inside a [`kill_on_overdraft`] scope and we haven't
-/// already started panicking for it, fire the kill panic with `bank_balance`
-/// as the payload.
+/// Fire the overdraft panic if the current task is inside a
+/// [`kill_on_overdraft`] scope and hasn't already panicked once. No-op
+/// otherwise — outside the scope or during drop-chain unwinding.
 #[cold]
 #[inline(never)]
 fn maybe_kill(bank_balance: isize) {
-    let _ = KILL_GUARD.try_with(|panicking| {
-        if !panicking.get() {
-            panicking.set(true);
+    let _ = DISALLOW_OVERDRAFT.try_with(|already_fired| {
+        if !already_fired.get() {
+            already_fired.set(true);
             std::panic::panic_any(OverdraftPanic { bank_balance });
         }
     });
@@ -293,12 +287,12 @@ impl AccountingAllocator<System> {
 unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
         // SAFETY: layout is forwarded unchanged.
-        let p = unsafe { self.inner.alloc(layout) };
-        if !p.is_null() {
+        let ptr = unsafe { self.inner.alloc(layout) };
+        if !ptr.is_null() {
             // Allocation debits the bank.
             track(-(layout.size() as isize));
         }
-        p
+        ptr
     }
 
     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
@@ -310,21 +304,21 @@ unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAllocator<A> {
 
     unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
         // SAFETY: layout is forwarded unchanged.
-        let p = unsafe { self.inner.alloc_zeroed(layout) };
-        if !p.is_null() {
+        let ptr = unsafe { self.inner.alloc_zeroed(layout) };
+        if !ptr.is_null() {
             track(-(layout.size() as isize));
         }
-        p
+        ptr
     }
 
     unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
         // SAFETY: caller upholds GlobalAlloc invariants; we forward unchanged.
-        let p = unsafe { self.inner.realloc(ptr, layout, new_size) };
-        if !p.is_null() {
+        let new_ptr = unsafe { self.inner.realloc(ptr, layout, new_size) };
+        if !new_ptr.is_null() {
             // Growth debits, shrink credits.
             track(layout.size() as isize - new_size as isize);
         }
-        p
+        new_ptr
     }
 }
 
@@ -349,7 +343,7 @@ mod tests {
         settle_thread_local();
         let before = bank_balance();
 
-        let v: Vec<u8> = vec![0u8; 8192];
+        let buf: Vec<u8> = vec![0u8; 8192];
         settle_thread_local();
         let mid = bank_balance();
         // Alloc debited the bank → mid should be at least 8192 below before.
@@ -358,7 +352,7 @@ mod tests {
             "alloc didn't debit: before={before} mid={mid}"
         );
 
-        drop(v);
+        drop(buf);
         settle_thread_local();
         let after = bank_balance();
         // Free credited the bank → after should be at least 8192 above mid.
@@ -374,10 +368,10 @@ mod tests {
         init_bank(1_000_000);
         // The bank drifts a little from this thread's own allocator activity
         // between init and the read, so we expect at-or-below the set value.
-        let v = bank_balance();
+        let bal = bank_balance();
         assert!(
-            (900_000..=1_000_000).contains(&v),
-            "init_bank didn't stick: v={v}"
+            (900_000..=1_000_000).contains(&bal),
+            "init_bank didn't stick: bal={bal}"
         );
     }
 
@@ -389,7 +383,7 @@ mod tests {
 
         // Alloc large enough to cross SETTLE_THRESHOLD in one shot — settle
         // drives balance from +1024 to roughly -(threshold+budget+overhead).
-        let _v: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 4096];
+        let _buf: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 4096];
         assert!(
             bank_balance() < 0,
             "bank should be negative after overdraft; bank={}",
@@ -405,7 +399,7 @@ mod tests {
         let result = std::panic::AssertUnwindSafe(kill_on_overdraft(async {
             // Big alloc crosses SETTLE_THRESHOLD → settle drives balance
             // negative → maybe_kill fires inside the settle closure → panic.
-            let _v: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 4096];
+            let _buf: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 4096];
             unreachable!("should have panicked");
         }))
         .catch_unwind()
@@ -428,7 +422,7 @@ mod tests {
         // No `kill_on_overdraft` wrapper -- overdraft should drive the bank
         // negative but allocations keep proceeding without panicking.
         init_bank(1024);
-        let _v: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 4096];
+        let _buf: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 4096];
         assert!(bank_balance() < 0, "should have overdrawn");
         // Reaching here without panic is the whole point.
     }
@@ -439,7 +433,7 @@ mod tests {
         settle_thread_local();
         let before_bank = bank_balance();
 
-        let v: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 1024];
+        let buf: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 1024];
         // Crossing the threshold auto-settles; bank should have dropped by
         // at least SETTLE_THRESHOLD without us calling settle_thread_local.
         let bank_after_alloc = bank_balance();
@@ -448,6 +442,6 @@ mod tests {
             "bank didn't auto-settle on threshold crossing: \
              before={before_bank} after_alloc={bank_after_alloc}"
         );
-        drop(v);
+        drop(buf);
     }
 }

From 69aaf1dc03c23002f50f662d1b147a0428fb37cd Mon Sep 17 00:00:00 2001
From: Brent Gardner <bgardner@squarelabs.net>
Date: Mon, 1 Jun 2026 12:27:39 -0600
Subject: [PATCH 14/17] fix wording

---
 datafusion/sqllogictest/bin/sqllogictests.rs  | 15 +--
 datafusion/sqllogictest/src/accounting.rs     | 99 +++++++++----------
 .../src/engines/datafusion_engine/runner.rs   |  8 +-
 datafusion/sqllogictest/src/lib.rs            |  2 +-
 4 files changed, 62 insertions(+), 62 deletions(-)

diff --git a/datafusion/sqllogictest/bin/sqllogictests.rs b/datafusion/sqllogictest/bin/sqllogictests.rs
index 43a4623fd8209..3fcf809061c3a 100644
--- a/datafusion/sqllogictest/bin/sqllogictests.rs
+++ b/datafusion/sqllogictest/bin/sqllogictests.rs
@@ -967,10 +967,11 @@ struct Options {
     #[clap(
         long,
         help = "Total memory budget for the whole process in MB (cgroup limit, \
-                container size, etc.). The allocator bank is set to this value, \
-                and panics any kill_on_overdraft-marked future as soon as real \
-                bytes allocated since startup exceed it. Requires the \
-                memory-accounting feature; ignored without it."
+                container size, etc.). Each per-file SLT context gets its own \
+                account in the allocator bank seeded with this budget; any \
+                kill_on_overdraft-marked future panics as soon as its account \
+                goes negative. Requires the memory-accounting feature; ignored \
+                without it."
     )]
     total_memory_mb: Option<usize>,
 
@@ -979,9 +980,9 @@ struct Options {
         default_value_t = 0.6,
         help = "DataFusion's share of --total-memory-mb. The per-file \
                 MemoryPool is sized at total * this fraction; the rest of the \
-                process budget covers hyper, tokio internals, gRPC, and other \
-                non-DF consumers. The bank catches violations of the overall \
-                budget regardless of which consumer caused the drift."
+                budget covers hyper, tokio internals, gRPC, and other non-DF \
+                consumers. Each file's account catches violations of its \
+                overall budget regardless of which consumer caused the drift."
     )]
     datafusion_memory_fraction: f64,
 }
diff --git a/datafusion/sqllogictest/src/accounting.rs b/datafusion/sqllogictest/src/accounting.rs
index 78aacec0c228d..384bbd30dc415 100644
--- a/datafusion/sqllogictest/src/accounting.rs
+++ b/datafusion/sqllogictest/src/accounting.rs
@@ -15,33 +15,31 @@
 // specific language governing permissions and limitations
 // under the License.
 
-//! Allocation-budget `GlobalAlloc` wrapper.
+//! Allocator-driven memory accounting with per-context budgets.
 //!
-//! The bank is a *budget*, not a counter. Allocations *debit* the bank,
-//! deallocations *credit* it. [`init_bank`] sets the starting budget.
-//! Going below zero is an overdraft — detectable via [`is_overdrawn`] or
-//! by reading [`bank_balance`] directly.
+//! The bank ([`ACCOUNTS`]) holds one [`AtomicIsize`] account per stamped
+//! `CONTEXT_ID`, each tracking its own remaining budget. Allocations debit
+//! the current thread's account, deallocations credit it; below zero is an
+//! overdraft. Threads with `CONTEXT_ID == 0` (main, the outer orchestration
+//! runtime, blocking-pool hosts) are untracked and skip the hot path.
 //!
-//! Every `alloc`/`realloc`/`dealloc` updates a thread-local `isize` (also
-//! a balance — debit on alloc, credit on free). When that local's magnitude
-//! crosses [`SETTLE_THRESHOLD`] it transfers into the global bank with one
-//! atomic op and the local zeros out. This amortizes cache-line contention
-//! across thousands of allocations.
+//! Per-alloc bookkeeping accumulates in a thread-local `LOCAL_BALANCE`
+//! drift counter; it settles into the account once `|drift|` crosses
+//! [`SETTLE_THRESHOLD`] (64 KB), amortizing the `RwLock` read + atomic
+//! op across thousands of allocations.
 //!
-//! [`bank_balance`] reads the global bank with a single relaxed atomic load.
-//! It lags reality by up to `SETTLE_THRESHOLD * num_threads` (≈1 MB on a
-//! 16-core box) — fine for catching multi-GB drift.
+//! [`account_balance`] reads the current thread's account; it lags reality
+//! by up to one threshold's worth of un-settled drift per thread.
 //!
 //! # Enforcement
 //!
-//! By default the allocator only *counts* — overdrafts set the sticky
-//! [`is_overdrawn`] flag but allocations continue. Code that wants the
-//! allocator to actively kill its in-flight work wraps its futures in
-//! [`kill_on_overdraft`]. While such a future is being polled, the very
-//! next debit after the bank goes negative panics that future with an
-//! [`OverdraftPanic`] payload. Other futures (system tasks, HTTP handlers)
-//! polled on the same worker thread are unaffected — they're not inside
-//! the scope, so `track()` skips the check.
+//! By default an overdraft just drives the account negative; allocations
+//! keep going. Wrap a future in [`kill_on_overdraft`] to make the allocator
+//! actively kill it on overdraft: the first settle that pushes the account
+//! below zero fires an [`OverdraftPanic`] on the polling thread. Drop
+//! chains during unwind don't re-panic (one-shot via [`DISALLOW_OVERDRAFT`]'s
+//! already-fired cell). Other futures on the same worker thread are
+//! unaffected — they're outside the scope, so `track` skips the check.
 //!
 //! Compiled in only when the `memory-accounting` feature is on.
 
@@ -150,8 +148,8 @@ tokio::task_local! {
 /// ```
 #[derive(Debug, Clone)]
 pub struct OverdraftPanic {
-    /// Bank balance at the moment the panic fired (negative — that's the point).
-    pub bank_balance: isize,
+    /// Account balance at the moment the panic fired (negative — that's the point).
+    pub account_balance: isize,
 }
 
 /// Wrap a future so the allocator may panic it on overdraft.
@@ -190,7 +188,7 @@ pub fn memory_tracker_limit() -> usize {
 }
 
 /// Current account balance. Negative = overdraft. `0` if untracked.
-pub fn bank_balance() -> isize {
+pub fn account_balance() -> isize {
     with_current_balance(|bal| bal.load(Ordering::Relaxed)).unwrap_or(0)
 }
 
@@ -254,11 +252,11 @@ fn track(delta: isize) {
 /// otherwise — outside the scope or during drop-chain unwinding.
 #[cold]
 #[inline(never)]
-fn maybe_kill(bank_balance: isize) {
+fn maybe_kill(account_balance: isize) {
     let _ = DISALLOW_OVERDRAFT.try_with(|already_fired| {
         if !already_fired.get() {
             already_fired.set(true);
-            std::panic::panic_any(OverdraftPanic { bank_balance });
+            std::panic::panic_any(OverdraftPanic { account_balance });
         }
     });
 }
@@ -338,15 +336,15 @@ mod tests {
     }
 
     #[test]
-    fn alloc_debits_and_free_credits_bank() {
+    fn alloc_debits_and_free_credits_account() {
         enter_fresh_context();
         settle_thread_local();
-        let before = bank_balance();
+        let before = account_balance();
 
         let buf: Vec<u8> = vec![0u8; 8192];
         settle_thread_local();
-        let mid = bank_balance();
-        // Alloc debited the bank → mid should be at least 8192 below before.
+        let mid = account_balance();
+        // Alloc debited the account → mid should be at least 8192 below before.
         assert!(
             before - mid >= 8192,
             "alloc didn't debit: before={before} mid={mid}"
@@ -354,8 +352,8 @@ mod tests {
 
         drop(buf);
         settle_thread_local();
-        let after = bank_balance();
-        // Free credited the bank → after should be at least 8192 above mid.
+        let after = account_balance();
+        // Free credited the account → after should be at least 8192 above mid.
         assert!(
             after - mid >= 8192,
             "free didn't credit: mid={mid} after={after}"
@@ -368,7 +366,7 @@ mod tests {
         init_bank(1_000_000);
         // The bank drifts a little from this thread's own allocator activity
         // between init and the read, so we expect at-or-below the set value.
-        let bal = bank_balance();
+        let bal = account_balance();
         assert!(
             (900_000..=1_000_000).contains(&bal),
             "init_bank didn't stick: bal={bal}"
@@ -376,18 +374,18 @@ mod tests {
     }
 
     #[test]
-    fn bank_goes_negative_on_overdraft() {
+    fn balance_goes_negative_on_overdraft() {
         enter_fresh_context();
         init_bank(1024);
-        assert!(bank_balance() > 0, "should start positive after init_bank");
+        assert!(account_balance() > 0, "should start positive after init_bank");
 
         // Alloc large enough to cross SETTLE_THRESHOLD in one shot — settle
         // drives balance from +1024 to roughly -(threshold+budget+overhead).
         let _buf: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 4096];
         assert!(
-            bank_balance() < 0,
-            "bank should be negative after overdraft; bank={}",
-            bank_balance()
+            account_balance() < 0,
+            "balance should be negative after overdraft; bal={}",
+            account_balance()
         );
     }
 
@@ -410,37 +408,38 @@ mod tests {
             .downcast_ref::<OverdraftPanic>()
             .expect("panic payload should be OverdraftPanic");
         assert!(
-            overdraft.bank_balance < 0,
+            overdraft.account_balance < 0,
             "payload should report negative balance; got {}",
-            overdraft.bank_balance
+            overdraft.account_balance
         );
     }
 
     #[tokio::test(flavor = "current_thread")]
     async fn no_kill_outside_scope() {
         enter_fresh_context();
-        // No `kill_on_overdraft` wrapper -- overdraft should drive the bank
+        // No `kill_on_overdraft` wrapper -- overdraft should drive the balance
         // negative but allocations keep proceeding without panicking.
         init_bank(1024);
         let _buf: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 4096];
-        assert!(bank_balance() < 0, "should have overdrawn");
+        assert!(account_balance() < 0, "should have overdrawn");
         // Reaching here without panic is the whole point.
     }
 
     #[test]
-    fn threshold_settlement_flushes_to_bank() {
+    fn threshold_settlement_flushes_to_account() {
         enter_fresh_context();
         settle_thread_local();
-        let before_bank = bank_balance();
+        let before = account_balance();
 
         let buf: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 1024];
-        // Crossing the threshold auto-settles; bank should have dropped by
-        // at least SETTLE_THRESHOLD without us calling settle_thread_local.
-        let bank_after_alloc = bank_balance();
+        // Crossing the threshold auto-settles; account balance should have
+        // dropped by at least SETTLE_THRESHOLD without us calling
+        // settle_thread_local.
+        let after_alloc = account_balance();
         assert!(
-            before_bank - bank_after_alloc >= SETTLE_THRESHOLD,
-            "bank didn't auto-settle on threshold crossing: \
-             before={before_bank} after_alloc={bank_after_alloc}"
+            before - after_alloc >= SETTLE_THRESHOLD,
+            "balance didn't auto-settle on threshold crossing: \
+             before={before} after_alloc={after_alloc}"
         );
         drop(buf);
     }
diff --git a/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs b/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs
index 82889d17c911a..7e84d0302f478 100644
--- a/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs
+++ b/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs
@@ -112,14 +112,14 @@ impl DataFusion {
                             / (1024 * 1024);
                         warn!(
                             "[{}] killed by allocator overdraft: \
-                             bank balance = {} bytes, df-pool reserved = {df_reserved_mb} MB; \
+                             account balance = {} bytes, df-pool reserved = {df_reserved_mb} MB; \
                              sql = {sql:?}",
                             self.relative_path.display(),
-                            od.bank_balance,
+                            od.account_balance,
                         );
                         Err(DFSqlLogicTestError::Other(format!(
-                            "allocator overdraft: bank balance at panic = {} bytes",
-                            od.bank_balance,
+                            "allocator overdraft: account balance at panic = {} bytes",
+                            od.account_balance,
                         )))
                     } else {
                         // Not our panic — re-raise so test runner sees it.
diff --git a/datafusion/sqllogictest/src/lib.rs b/datafusion/sqllogictest/src/lib.rs
index 00735dd92f4b8..57537567b3568 100644
--- a/datafusion/sqllogictest/src/lib.rs
+++ b/datafusion/sqllogictest/src/lib.rs
@@ -33,7 +33,7 @@ mod test_file;
 
 #[cfg(feature = "memory-accounting")]
 pub use accounting::{
-    AccountingAllocator, OverdraftPanic, bank_balance, current_context_id, init_bank,
+    AccountingAllocator, OverdraftPanic, account_balance, current_context_id, init_bank,
     kill_on_overdraft, local_balance, memory_tracker_limit, next_context_id,
     set_memory_tracker_limit, set_thread_context_id, settle_thread_local,
 };

From c09dd29c6d4d9b5a0d411ba5b8b21d459d65990a Mon Sep 17 00:00:00 2001
From: Brent Gardner <bgardner@squarelabs.net>
Date: Mon, 1 Jun 2026 12:36:18 -0600
Subject: [PATCH 15/17] clean up naming and shape after the per-account
 refactor

The Step-3 split (one bank with N accounts, one per stamped context-id)
left a lot of vestigial vocabulary and shape from the old "one global
bank" model. Tightening it now so the code reads in one voice.

Naming:
  - bank_balance -> account_balance (fn, OverdraftPanic field, locals,
    log strings). The bank is the system of accounts; only an account
    has a balance.
  - KILL_GUARD -> DISALLOW_OVERDRAFT, panicking -> already_fired. The
    task_local names the policy the scope enforces, not the mechanism;
    matches the kill_on_overdraft / OverdraftPanic vocabulary.
  - init_bank -> set_default_budget + set_account_balance. The old
    name conflated two unrelated concerns (bank-level template budget
    vs current-thread account reset). Two single-purpose fns let the
    SLT runner call just the bank-level one at startup and tests seed
    a freshly-stamped account explicitly.
  - Single-letter vars killed throughout: |op|, |future|, |poison|,
    drift, new_bal, ptr/new_ptr, buf. Closures over the LOCAL_BALANCE
    cell use |loc_bal|; closures over an account's &AtomicIsize use
    |bal|; both consistent.

Shape:
  - track() flipped to the early-return railroad. The 99% case
    (drift fits under threshold) is the first three lines inside the
    closure and returns; the settle path runs flat. Max indent inside
    the function body dropped from 5 to 2.
  - maybe_kill takes the new balance from the caller; no more
    redundant lookup that would have required re-acquiring the lock.
  - The kill panic fires from outside the read closure, so the
    payload allocation can't recurse-deadlock on the same RwLock.

Behavior unchanged: tests green, SLT integration green except for the
pre-existing information_schema.slt MemoryPool-display diff that's
unrelated to bank state.
---
 datafusion/sqllogictest/bin/sqllogictests.rs  |  2 +-
 datafusion/sqllogictest/src/accounting.rs     | 37 +++++++++++--------
 .../src/engines/datafusion_engine/runner.rs   | 10 ++---
 datafusion/sqllogictest/src/lib.rs            |  5 ++-
 4 files changed, 29 insertions(+), 25 deletions(-)

diff --git a/datafusion/sqllogictest/bin/sqllogictests.rs b/datafusion/sqllogictest/bin/sqllogictests.rs
index 3fcf809061c3a..e17a3fbd7c9a2 100644
--- a/datafusion/sqllogictest/bin/sqllogictests.rs
+++ b/datafusion/sqllogictest/bin/sqllogictests.rs
@@ -152,7 +152,7 @@ async fn run_tests() -> Result<()> {
         // process allocations exceed `total_bytes`, regardless of whether
         // DF's voluntary tracker noticed.
         datafusion_sqllogictest::set_memory_tracker_limit(df_bytes);
-        datafusion_sqllogictest::init_bank(total_bytes as isize);
+        datafusion_sqllogictest::set_default_budget(total_bytes as isize);
         log::info!(
             "memory-accounting on: total={total_mb} MB, datafusion={} MB \
              (fraction {:.2}), non-datafusion budget={} MB",
diff --git a/datafusion/sqllogictest/src/accounting.rs b/datafusion/sqllogictest/src/accounting.rs
index 384bbd30dc415..90590b39d3b89 100644
--- a/datafusion/sqllogictest/src/accounting.rs
+++ b/datafusion/sqllogictest/src/accounting.rs
@@ -60,8 +60,8 @@ const SETTLE_THRESHOLD: isize = 64 * 1024;
 /// gets an entry — that's the "untracked thread" marker.
 static ACCOUNTS: OnceLock<RwLock<HashMap<usize, AtomicIsize>>> = OnceLock::new();
 
-/// Starting budget for any new account, set by [`init_bank`] and inherited
-/// by per-file SLT contexts spawned after.
+/// Starting budget for any new account, set by [`set_default_budget`] and
+/// inherited by per-file SLT contexts spawned after.
 static DEFAULT_BUDGET: AtomicIsize = AtomicIsize::new(0);
 
 fn accounts() -> &'static RwLock<HashMap<usize, AtomicIsize>> {
@@ -117,7 +117,9 @@ pub fn set_thread_context_id(id: usize) {
     // read-lock the map we're holding for writing on the same thread.
     {
         let accounts_lock = accounts();
-        let mut accounts = accounts_lock.write().unwrap_or_else(|poison| poison.into_inner());
+        let mut accounts = accounts_lock
+            .write()
+            .unwrap_or_else(|poison| poison.into_inner());
         accounts
             .entry(id)
             .or_insert_with(|| AtomicIsize::new(DEFAULT_BUDGET.load(Ordering::Relaxed)));
@@ -164,10 +166,15 @@ pub async fn kill_on_overdraft<F: Future>(future: F) -> F::Output {
     DISALLOW_OVERDRAFT.scope(Cell::new(false), future).await
 }
 
-/// Set the default budget for accounts created later, and reset the
-/// current thread's account (if any) to `value`.
-pub fn init_bank(value: isize) {
+/// Set the default budget new accounts will be created with. Existing
+/// accounts are untouched.
+pub fn set_default_budget(value: isize) {
     DEFAULT_BUDGET.store(value, Ordering::Relaxed);
+}
+
+/// Set the current thread's account balance to `value`. No-op on untracked
+/// threads (`CONTEXT_ID == 0`).
+pub fn set_account_balance(value: isize) {
     let _ = with_current_balance(|bal| bal.store(value, Ordering::Relaxed));
 }
 
@@ -361,23 +368,23 @@ mod tests {
     }
 
     #[test]
-    fn init_bank_sets_value() {
+    fn set_account_balance_sticks() {
         enter_fresh_context();
-        init_bank(1_000_000);
-        // The bank drifts a little from this thread's own allocator activity
-        // between init and the read, so we expect at-or-below the set value.
+        set_account_balance(1_000_000);
+        // Balance drifts a little from this thread's own allocator activity
+        // between the set and the read, so we expect at-or-below the set value.
         let bal = account_balance();
         assert!(
             (900_000..=1_000_000).contains(&bal),
-            "init_bank didn't stick: bal={bal}"
+            "set_account_balance didn't stick: bal={bal}"
         );
     }
 
     #[test]
     fn balance_goes_negative_on_overdraft() {
         enter_fresh_context();
-        init_bank(1024);
-        assert!(account_balance() > 0, "should start positive after init_bank");
+        set_account_balance(1024);
+        assert!(account_balance() > 0, "should start positive");
 
         // Alloc large enough to cross SETTLE_THRESHOLD in one shot — settle
         // drives balance from +1024 to roughly -(threshold+budget+overhead).
@@ -392,7 +399,7 @@ mod tests {
     #[tokio::test(flavor = "current_thread")]
     async fn kill_on_overdraft_panics_inside_scope() {
         enter_fresh_context();
-        init_bank(1024);
+        set_account_balance(1024);
 
         let result = std::panic::AssertUnwindSafe(kill_on_overdraft(async {
             // Big alloc crosses SETTLE_THRESHOLD → settle drives balance
@@ -417,9 +424,9 @@ mod tests {
     #[tokio::test(flavor = "current_thread")]
     async fn no_kill_outside_scope() {
         enter_fresh_context();
+        set_account_balance(1024);
         // No `kill_on_overdraft` wrapper -- overdraft should drive the balance
         // negative but allocations keep proceeding without panicking.
-        init_bank(1024);
         let _buf: Vec<u8> = vec![0u8; SETTLE_THRESHOLD as usize + 4096];
         assert!(account_balance() < 0, "should have overdrawn");
         // Reaching here without panic is the whole point.
diff --git a/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs b/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs
index 7e84d0302f478..6e737f888bca8 100644
--- a/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs
+++ b/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs
@@ -103,13 +103,9 @@ impl DataFusion {
                 Ok(r) => r,
                 Err(payload) => {
                     if let Some(od) = payload.downcast_ref::<OverdraftPanic>() {
-                        let df_reserved_mb = (self
-                            .ctx
-                            .runtime_env()
-                            .memory_pool
-                            .reserved()
-                            as u64)
-                            / (1024 * 1024);
+                        let df_reserved_mb =
+                            (self.ctx.runtime_env().memory_pool.reserved() as u64)
+                                / (1024 * 1024);
                         warn!(
                             "[{}] killed by allocator overdraft: \
                              account balance = {} bytes, df-pool reserved = {df_reserved_mb} MB; \
diff --git a/datafusion/sqllogictest/src/lib.rs b/datafusion/sqllogictest/src/lib.rs
index 57537567b3568..b056e62fd2121 100644
--- a/datafusion/sqllogictest/src/lib.rs
+++ b/datafusion/sqllogictest/src/lib.rs
@@ -33,9 +33,10 @@ mod test_file;
 
 #[cfg(feature = "memory-accounting")]
 pub use accounting::{
-    AccountingAllocator, OverdraftPanic, account_balance, current_context_id, init_bank,
+    AccountingAllocator, OverdraftPanic, account_balance, current_context_id,
     kill_on_overdraft, local_balance, memory_tracker_limit, next_context_id,
-    set_memory_tracker_limit, set_thread_context_id, settle_thread_local,
+    set_account_balance, set_default_budget, set_memory_tracker_limit,
+    set_thread_context_id, settle_thread_local,
 };
 
 pub use engines::CurrentlyExecutingSqlTracker;

From de927c84141b806648843dd12a7619a14c583948 Mon Sep 17 00:00:00 2001
From: Brent Gardner <bgardner@squarelabs.net>
Date: Mon, 1 Jun 2026 14:15:06 -0600
Subject: [PATCH 16/17] Per threadpool tracking

---
 .github/workflows/rust.yml                    |   2 +-
 datafusion/execution/src/memory_pool/mod.rs   |  12 +-
 datafusion/execution/src/memory_pool/pool.rs  |  29 ++++-
 datafusion/execution/src/runtime_env.rs       |  15 ++-
 datafusion/sqllogictest/bin/sqllogictests.rs  |  54 +++-----
 .../sqllogictest/src/accounting_pool.rs       | 122 ++++++++++++++++++
 datafusion/sqllogictest/src/lib.rs            |   4 +
 datafusion/sqllogictest/src/test_context.rs   |  20 ++-
 8 files changed, 207 insertions(+), 51 deletions(-)
 create mode 100644 datafusion/sqllogictest/src/accounting_pool.rs

diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml
index ac23a9775b8dd..f167117d5d146 100644
--- a/.github/workflows/rust.yml
+++ b/.github/workflows/rust.yml
@@ -471,7 +471,7 @@ jobs:
           export RUST_MIN_STACK=20971520
           export TPCH_DATA=`realpath datafusion/sqllogictest/test_files/tpch/data`
           cargo test plan_q --package datafusion-benchmarks --profile ci --features=ci -- --test-threads=1
-          INCLUDE_TPCH=true cargo test --features backtrace,parquet_encryption,substrait,memory-accounting --profile ci --package datafusion-sqllogictest --test sqllogictests -- --total-memory-mb 16384 --datafusion-memory-fraction 0.6
+          INCLUDE_TPCH=true cargo test --features backtrace,parquet_encryption,substrait,memory-accounting --profile ci --package datafusion-sqllogictest --test sqllogictests -- --default-pool-size-mb 16384
       - name: Verify Working Directory Clean
         run: git diff --exit-code
 
diff --git a/datafusion/execution/src/memory_pool/mod.rs b/datafusion/execution/src/memory_pool/mod.rs
index 2b36ee7f40add..e50f72632b3f2 100644
--- a/datafusion/execution/src/memory_pool/mod.rs
+++ b/datafusion/execution/src/memory_pool/mod.rs
@@ -18,7 +18,7 @@
 //! [`MemoryPool`] for memory management during query execution, [`proxy`] for
 //! help with allocation accounting.
 
-use datafusion_common::{Result, internal_datafusion_err};
+use datafusion_common::{Result, internal_datafusion_err, not_impl_err};
 use std::any::Any;
 use std::fmt::Display;
 use std::hash::{Hash, Hasher};
@@ -223,6 +223,16 @@ pub trait MemoryPool: Any + Send + Sync + std::fmt::Debug + Display {
     fn memory_limit(&self) -> MemoryLimit {
         MemoryLimit::Unknown
     }
+
+    /// Attempt to update this pool's limit in place to `new_limit` bytes.
+    ///
+    /// Default impl returns `Err`. Callers that route through
+    /// [`crate::runtime_env::RuntimeEnvBuilder::with_memory_limit`] fall
+    /// back to replacing the pool wholesale on `Err`, preserving historical
+    /// behavior for pools that can't be resized in place.
+    fn try_resize(&self, _new_limit: usize) -> Result<()> {
+        not_impl_err!("{} does not support resize", self.name())
+    }
 }
 
 impl dyn MemoryPool {
diff --git a/datafusion/execution/src/memory_pool/pool.rs b/datafusion/execution/src/memory_pool/pool.rs
index 52b601d5cd78b..ce8b76d63a610 100644
--- a/datafusion/execution/src/memory_pool/pool.rs
+++ b/datafusion/execution/src/memory_pool/pool.rs
@@ -73,9 +73,15 @@ impl Display for UnboundedMemoryPool {
 /// This pool works well for queries that do not need to spill or have
 /// a single spillable operator. See [`FairSpillPool`] if there are
 /// multiple spillable operators that all will spill.
+///
+/// Supports [`MemoryPool::try_resize`] for in-place limit adjustment, so
+/// callers routing through
+/// [`RuntimeEnvBuilder::with_memory_limit`](crate::runtime_env::RuntimeEnvBuilder::with_memory_limit)
+/// can keep the existing pool (and any wrappers around it) rather than
+/// replacing it on every change.
 #[derive(Debug)]
 pub struct GreedyMemoryPool {
-    pool_size: usize,
+    pool_size: AtomicUsize,
     used: AtomicUsize,
 }
 
@@ -84,7 +90,7 @@ impl GreedyMemoryPool {
     pub fn new(pool_size: usize) -> Self {
         debug!("Created new GreedyMemoryPool(pool_size={pool_size})");
         Self {
-            pool_size,
+            pool_size: AtomicUsize::new(pool_size),
             used: AtomicUsize::new(0),
         }
     }
@@ -104,16 +110,17 @@ impl MemoryPool for GreedyMemoryPool {
     }
 
     fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> {
+        let pool_size = self.pool_size.load(Ordering::Relaxed);
         self.used
             .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |used| {
                 let new_used = used + additional;
-                (new_used <= self.pool_size).then_some(new_used)
+                (new_used <= pool_size).then_some(new_used)
             })
             .map_err(|used| {
                 insufficient_capacity_err(
                     reservation,
                     additional,
-                    self.pool_size.saturating_sub(used),
+                    pool_size.saturating_sub(used),
                     self,
                 )
             })?;
@@ -125,19 +132,25 @@ impl MemoryPool for GreedyMemoryPool {
     }
 
     fn memory_limit(&self) -> MemoryLimit {
-        MemoryLimit::Finite(self.pool_size)
+        MemoryLimit::Finite(self.pool_size.load(Ordering::Relaxed))
+    }
+
+    fn try_resize(&self, new_limit: usize) -> Result<()> {
+        self.pool_size.store(new_limit, Ordering::Relaxed);
+        Ok(())
     }
 }
 
 impl Display for GreedyMemoryPool {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         let used = self.used.load(Ordering::Relaxed);
+        let pool_size = self.pool_size.load(Ordering::Relaxed);
         write!(
             f,
             "{}(used: {}, pool_size: {})",
             &self.name(),
             human_readable_size(used),
-            human_readable_size(self.pool_size)
+            human_readable_size(pool_size)
         )
     }
 }
@@ -600,6 +613,10 @@ impl<I: MemoryPool> MemoryPool for TrackConsumersPool<I> {
     fn memory_limit(&self) -> MemoryLimit {
         self.inner.memory_limit()
     }
+
+    fn try_resize(&self, new_limit: usize) -> Result<()> {
+        self.inner.try_resize(new_limit)
+    }
 }
 
 fn provide_top_memory_consumers_to_error_msg(
diff --git a/datafusion/execution/src/runtime_env.rs b/datafusion/execution/src/runtime_env.rs
index 5b90f28a141ef..b8573f22b636a 100644
--- a/datafusion/execution/src/runtime_env.rs
+++ b/datafusion/execution/src/runtime_env.rs
@@ -409,12 +409,23 @@ impl RuntimeEnvBuilder {
     /// Specify the total memory to use while running the DataFusion
     /// plan to `max_memory * memory_fraction` in bytes.
     ///
-    /// This defaults to using [`GreedyMemoryPool`] wrapped in the
-    /// [`TrackConsumersPool`] with a maximum of 5 consumers.
+    /// If a memory pool is already configured on this builder, this first
+    /// attempts to resize it in place via [`MemoryPool::try_resize`]. Pools
+    /// that support resize (e.g. [`GreedyMemoryPool`]) keep their identity
+    /// — useful for any wrapper that needs to observe limit changes (e.g.
+    /// to retune external accounting). Pools whose [`MemoryPool::try_resize`]
+    /// returns `Err` (the default) fall back to wholesale replacement
+    /// with a [`TrackConsumersPool`]-wrapped [`GreedyMemoryPool`] (top 5
+    /// consumers), preserving the historical behavior.
     ///
     /// Note DataFusion does not yet respect this limit in all cases.
     pub fn with_memory_limit(self, max_memory: usize, memory_fraction: f64) -> Self {
         let pool_size = (max_memory as f64 * memory_fraction) as usize;
+        if let Some(existing) = &self.memory_pool {
+            if existing.try_resize(pool_size).is_ok() {
+                return self;
+            }
+        }
         self.with_memory_pool(Arc::new(TrackConsumersPool::new(
             GreedyMemoryPool::new(pool_size),
             NonZeroUsize::new(5).unwrap(),
diff --git a/datafusion/sqllogictest/bin/sqllogictests.rs b/datafusion/sqllogictest/bin/sqllogictests.rs
index e17a3fbd7c9a2..8e6b397750fc5 100644
--- a/datafusion/sqllogictest/bin/sqllogictests.rs
+++ b/datafusion/sqllogictest/bin/sqllogictests.rs
@@ -143,23 +143,16 @@ async fn run_tests() -> Result<()> {
     options.warn_on_ignored();
 
     #[cfg(feature = "memory-accounting")]
-    if let Some(total_mb) = options.total_memory_mb {
-        let total_bytes = total_mb.saturating_mul(1024 * 1024);
-        let df_bytes = (total_bytes as f64 * options.datafusion_memory_fraction) as usize;
-        // The DataFusion MemoryPool gets its slice of the total budget; the
-        // allocator bank tracks the whole pie (DF + everything else in the
-        // process). A kill_on_overdraft-marked future panics when total
-        // process allocations exceed `total_bytes`, regardless of whether
-        // DF's voluntary tracker noticed.
-        datafusion_sqllogictest::set_memory_tracker_limit(df_bytes);
-        datafusion_sqllogictest::set_default_budget(total_bytes as isize);
-        log::info!(
-            "memory-accounting on: total={total_mb} MB, datafusion={} MB \
-             (fraction {:.2}), non-datafusion budget={} MB",
-            df_bytes / (1024 * 1024),
-            options.datafusion_memory_fraction,
-            (total_bytes - df_bytes) / (1024 * 1024)
-        );
+    if let Some(pool_mb) = options.default_pool_size_mb {
+        let pool_bytes = pool_mb.saturating_mul(1024 * 1024);
+        // Same value drives the inner MemoryPool's size and the bank's
+        // default budget. The wrapper renders this value as `unlimited` in
+        // `SHOW ALL` (sentinel for "no SET has happened"); once a test
+        // calls `SET datafusion.runtime.memory_limit`, the wrapper retunes
+        // the bank to that limit + 10% headroom.
+        datafusion_sqllogictest::set_memory_tracker_limit(pool_bytes);
+        datafusion_sqllogictest::set_default_budget(pool_bytes as isize);
+        log::info!("memory-accounting on: default pool size = {pool_mb} MB");
     }
 
     // Print parallelism info for debugging CI performance
@@ -966,25 +959,16 @@ struct Options {
 
     #[clap(
         long,
-        help = "Total memory budget for the whole process in MB (cgroup limit, \
-                container size, etc.). Each per-file SLT context gets its own \
-                account in the allocator bank seeded with this budget; any \
-                kill_on_overdraft-marked future panics as soon as its account \
-                goes negative. Requires the memory-accounting feature; ignored \
-                without it."
-    )]
-    total_memory_mb: Option<usize>,
-
-    #[clap(
-        long,
-        default_value_t = 0.6,
-        help = "DataFusion's share of --total-memory-mb. The per-file \
-                MemoryPool is sized at total * this fraction; the rest of the \
-                budget covers hyper, tokio internals, gRPC, and other non-DF \
-                consumers. Each file's account catches violations of its \
-                overall budget regardless of which consumer caused the drift."
+        help = "Default MemoryPool size in MB for each per-file SLT context. \
+                The pool is wrapped in AccountingMemoryPool, which doubles \
+                this value as the 'no SET has happened yet' sentinel — until \
+                an SLT calls `SET datafusion.runtime.memory_limit`, SHOW ALL \
+                renders the limit as 'unlimited' and the allocator bank \
+                stays loose. Once a test SETs a limit, the bank tightens to \
+                that limit + 10% headroom. Requires the memory-accounting \
+                feature; ignored without it."
     )]
-    datafusion_memory_fraction: f64,
+    default_pool_size_mb: Option<usize>,
 }
 
 impl Options {
diff --git a/datafusion/sqllogictest/src/accounting_pool.rs b/datafusion/sqllogictest/src/accounting_pool.rs
new file mode 100644
index 0000000000000..5a1970d78cd4a
--- /dev/null
+++ b/datafusion/sqllogictest/src/accounting_pool.rs
@@ -0,0 +1,122 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`AccountingMemoryPool`] bridges DataFusion's voluntary memory tracking
+//! to the allocator-level bank in [`crate::accounting`].
+//!
+//! It wraps any [`MemoryPool`] and re-tunes the current thread's bank
+//! account whenever the pool's limit changes (via [`MemoryPool::try_resize`],
+//! which `RuntimeEnvBuilder::with_memory_limit` triggers on `SET
+//! datafusion.runtime.memory_limit = '…'`).
+//!
+//! Each retune sets the bank to `new_limit * 1.10` — a fixed 10% headroom
+//! over what DataFusion thinks it's allowed to use, so a query that
+//! actually allocates >10% beyond its declared limit panics with an
+//! `OverdraftPanic`. That gap *is* the bug we're hunting: DataFusion's
+//! voluntary tracker saying one thing while the allocator says another.
+
+use crate::set_account_balance;
+use datafusion::common::Result;
+use datafusion::execution::memory_pool::{
+    MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation,
+};
+use std::fmt::{self, Display, Formatter};
+use std::sync::Arc;
+
+/// 10% headroom over the pool's declared limit. If actual allocations exceed
+/// this, DataFusion's voluntary tracking is lying.
+const HEADROOM_FACTOR: f64 = 1.10;
+
+pub struct AccountingMemoryPool {
+    inner: Arc<dyn MemoryPool>,
+    /// The operator-configured default pool size, used as a "no SET has
+    /// happened yet" sentinel by [`Self::memory_limit`].
+    default_size: usize,
+}
+
+impl AccountingMemoryPool {
+    pub fn new(inner: Arc<dyn MemoryPool>, default_size: usize) -> Self {
+        Self {
+            inner,
+            default_size,
+        }
+    }
+}
+
+impl fmt::Debug for AccountingMemoryPool {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        f.debug_struct("AccountingMemoryPool")
+            .field("inner", &self.inner)
+            .field("default_size", &self.default_size)
+            .finish()
+    }
+}
+
+impl Display for AccountingMemoryPool {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "accounting({})", self.inner)
+    }
+}
+
+impl MemoryPool for AccountingMemoryPool {
+    fn name(&self) -> &str {
+        "accounting"
+    }
+
+    fn register(&self, consumer: &MemoryConsumer) {
+        self.inner.register(consumer)
+    }
+
+    fn unregister(&self, consumer: &MemoryConsumer) {
+        self.inner.unregister(consumer)
+    }
+
+    fn grow(&self, reservation: &MemoryReservation, additional: usize) {
+        self.inner.grow(reservation, additional)
+    }
+
+    fn shrink(&self, reservation: &MemoryReservation, shrink: usize) {
+        self.inner.shrink(reservation, shrink)
+    }
+
+    fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> {
+        self.inner.try_grow(reservation, additional)
+    }
+
+    fn reserved(&self) -> usize {
+        self.inner.reserved()
+    }
+
+    fn memory_limit(&self) -> MemoryLimit {
+        // HACK: When the inner pool still reports the operator-configured
+        // default, no `SET datafusion.runtime.memory_limit` has happened —
+        // render as `Infinite` so `information_schema.slt`'s `SHOW ALL`
+        // expectation of `unlimited` for an un-SET context stays satisfied.
+        // Once a SET fires, `try_resize` mutates the inner pool to some
+        // other value and we report the real limit.
+        match self.inner.memory_limit() {
+            MemoryLimit::Finite(n) if n == self.default_size => MemoryLimit::Infinite,
+            other => other,
+        }
+    }
+
+    fn try_resize(&self, new_limit: usize) -> Result<()> {
+        self.inner.try_resize(new_limit)?;
+        set_account_balance((new_limit as f64 * HEADROOM_FACTOR) as isize);
+        Ok(())
+    }
+}
diff --git a/datafusion/sqllogictest/src/lib.rs b/datafusion/sqllogictest/src/lib.rs
index b056e62fd2121..206e5aadb8c5a 100644
--- a/datafusion/sqllogictest/src/lib.rs
+++ b/datafusion/sqllogictest/src/lib.rs
@@ -28,6 +28,8 @@
 
 #[cfg(feature = "memory-accounting")]
 mod accounting;
+#[cfg(feature = "memory-accounting")]
+mod accounting_pool;
 mod engines;
 mod test_file;
 
@@ -38,6 +40,8 @@ pub use accounting::{
     set_account_balance, set_default_budget, set_memory_tracker_limit,
     set_thread_context_id, settle_thread_local,
 };
+#[cfg(feature = "memory-accounting")]
+pub use accounting_pool::AccountingMemoryPool;
 
 pub use engines::CurrentlyExecutingSqlTracker;
 pub use engines::DFColumnType;
diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs
index 0994f0e3dcdbd..a1ed683088acd 100644
--- a/datafusion/sqllogictest/src/test_context.rs
+++ b/datafusion/sqllogictest/src/test_context.rs
@@ -95,19 +95,27 @@ impl TypePlanner for SqlLogicTestTypePlanner {
 }
 
 /// Construct the per-file `RuntimeEnv`. With the `memory-accounting` feature
-/// on and a non-zero `memory_tracker_limit()` configured, this builds a
-/// finite-pool runtime so DataFusion's voluntary tracker can be observed
-/// against the allocator's ground-truth bank. Otherwise falls back to the
-/// historical `RuntimeEnv::default()` (unbounded pool).
+/// on and a non-zero `memory_tracker_limit()` configured, this wraps the
+/// usual `TrackConsumersPool(GreedyMemoryPool)` in an `AccountingMemoryPool`
+/// so the allocator-level bank retunes on every `SET datafusion.runtime.
+/// memory_limit`. Otherwise falls back to the historical default.
 fn build_runtime_env() -> RuntimeEnv {
     #[cfg(feature = "memory-accounting")]
     {
+        use datafusion::execution::memory_pool::{GreedyMemoryPool, TrackConsumersPool};
+        use std::num::NonZeroUsize;
+
         let limit = crate::memory_tracker_limit();
         if limit > 0 {
+            let tracked = TrackConsumersPool::new(
+                GreedyMemoryPool::new(limit),
+                NonZeroUsize::new(5).unwrap(),
+            );
+            let wrapped = crate::AccountingMemoryPool::new(Arc::new(tracked), limit);
             return RuntimeEnvBuilder::new()
-                .with_memory_limit(limit, 1.0)
+                .with_memory_pool(Arc::new(wrapped))
                 .build()
-                .expect("RuntimeEnvBuilder::build with memory limit");
+                .expect("RuntimeEnvBuilder::build with accounting pool");
         }
     }
     RuntimeEnv::default()

From 706f3c7499ee527e11f1748e4284cb5ff8bce1c8 Mon Sep 17 00:00:00 2001
From: Brent Gardner <bgardner@squarelabs.net>
Date: Mon, 1 Jun 2026 14:19:07 -0600
Subject: [PATCH 17/17] doc updates

---
 datafusion/sqllogictest/README.md        | 30 ++++++++++++++++++++++++
 docs/source/contributor-guide/testing.md | 12 ++++++++++
 2 files changed, 42 insertions(+)

diff --git a/datafusion/sqllogictest/README.md b/datafusion/sqllogictest/README.md
index f0a54cf978fbf..a3c48cdde883b 100644
--- a/datafusion/sqllogictest/README.md
+++ b/datafusion/sqllogictest/README.md
@@ -360,6 +360,36 @@ For focusing on one specific failing test, a file:line filter can be used:
 cargo test --test sqllogictests -- --substrait-round-trip binary.slt:23
 ```
 
+## Running tests: allocator-level memory accounting
+
+Build with `--features memory-accounting` to install a global allocator
+wrapper that tracks actual bytes allocated per SLT file and reconciles them
+against DataFusion's voluntary `MemoryPool` tracking. The point isn't to
+enforce a process-wide budget — it's to catch DataFusion lying about how
+much memory it's using. If `MemoryPool` reports 1 MB while the allocator
+sees 100 MB go by, *that gap is the bug*.
+
+```shell
+cargo test --features memory-accounting --test sqllogictests -- \
+    --default-pool-size-mb 16384
+```
+
+`--default-pool-size-mb` seeds each per-file SLT context's MemoryPool with
+the given size in MB and arms the bank as a no-op until a test opts in.
+
+**Opting an individual test in.** Add `SET datafusion.runtime.memory_limit
+= 'N'` at the top of the `.slt`. The wrapping `AccountingMemoryPool` then
+tightens its allocator-level bank to `N * 1.10` (10% headroom). If the test
+allocates more than that — including bytes DataFusion's tracker didn't see
+— the test panics with an `OverdraftPanic` reporting the actual balance at
+panic time. SLTs without a `SET` of `memory_limit` see no change in
+behavior; the bank stays loose and `SHOW ALL` continues to render the limit
+as `unlimited`.
+
+Inside the runner each file gets its own multi-thread Tokio runtime so
+context-ids stamped onto worker threads stay stable for the allocator
+hook, and per-file accounts in the bank are isolated from each other.
+
 ## `.slt` file format
 
 [`sqllogictest`] was originally written for SQLite to verify the
diff --git a/docs/source/contributor-guide/testing.md b/docs/source/contributor-guide/testing.md
index 3b644f610b90e..3e44e3aabaeef 100644
--- a/docs/source/contributor-guide/testing.md
+++ b/docs/source/contributor-guide/testing.md
@@ -113,6 +113,18 @@ Like similar systems such as [DuckDB](https://duckdb.org/dev/testing), DataFusio
 
 DataFusion has integrated [sqlite's test suite](https://sqlite.org/sqllogictest/doc/trunk/about.wiki) as a supplemental test suite that is run whenever a PR is merged into DataFusion. To run it manually please refer to the [README](https://github.com/apache/datafusion/blob/main/datafusion/sqllogictest/README.md#running-tests-sqlite) file for instructions.
 
+### Allocator-level memory accounting (`--features memory-accounting`)
+
+For tests that need to verify DataFusion's voluntary memory tracking
+matches actual heap usage, the `sqllogictest` runner ships an optional
+`memory-accounting` feature that installs a global allocator wrapper.
+Adding `SET datafusion.runtime.memory_limit = 'N'` at the top of an
+`.slt` file opts that file into allocator-vs-`MemoryPool` reconciliation
+with 10% headroom — any divergence panics the test with an
+`OverdraftPanic` reporting the actual allocator balance. See
+[the sqllogictest README](https://github.com/apache/datafusion/blob/main/datafusion/sqllogictest/README.md#running-tests-allocator-level-memory-accounting)
+for the runner flag and the full mechanism.
+
 ## Snapshot testing (`cargo insta`)
 
 [Insta](https://github.com/mitsuhiko/insta) is used for snapshot testing. Snapshots are generated