Track "RSS" vs MemoryPool suring SLTs to prevent OOMs

.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
 

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",

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;
@@ -137,6 +142,26 @@ 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::init_bank(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)
+        );
+    }
+
     // Print parallelism info for debugging CI performance
     eprintln!(
         "Running with {} test threads (available parallelism: {})",
@@ -910,6 +935,27 @@ struct Options {
         default_value_t = ColorChoice::Auto
     )]
     color: ColorChoice,
+
+    #[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."
+    )]
+    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 \
+                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."
+    )]
+    datafusion_memory_fraction: f64,
 }
 
 impl Options {