From 04e87feb227843678d1076f67b44e9992dd80983 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 13 Jul 2026 14:20:47 -0600 Subject: [PATCH 1/2] feat: treat ResourcesExhausted as a retriable task failure A task that exhausts its memory budget currently fails the entire job on the first occurrence: `From for FailedTask` maps everything that is not an IO error to `FailedReason::ExecutionError`, and the scheduler fails the stage outright on that reason with no retry. A memory exhaustion is often transient. The retried task may land on a less-loaded executor, or run once its peers on the same executor have drained. Classify it as retriable so the scheduler reschedules it, bounded by --task-max-failures, rather than failing the job outright. Match the error through DataFusionError::find_root(), because DataFusion routinely wraps errors in Context, and in Shared when propagating one stream error to several output partitions (RepartitionExec). A bare match on the outer error would miss most real cases. The unpartitioned shuffle-write path flattened the error with `format!("{e:?}")` before it reached the classifier, which destroyed the type and made the new arm unreachable for exactly the stages most likely to exhaust memory: the final stage, the broadcast-join build side, and the CoalescePartitions and SortPreservingMerge input stages. Preserve the DataFusionError instead of formatting it. --- ballista/core/proto/ballista.proto | 6 + ballista/core/src/error.rs | 163 +++++++++++++++++- .../src/execution_plans/shuffle_writer.rs | 148 +++++++++++++++- ballista/core/src/serde/generated/ballista.rs | 8 +- ballista/scheduler/src/api/handlers.rs | 4 +- 5 files changed, 325 insertions(+), 4 deletions(-) diff --git a/ballista/core/proto/ballista.proto b/ballista/core/proto/ballista.proto index 41c946365f..8ecde3a01e 100644 --- a/ballista/core/proto/ballista.proto +++ b/ballista/core/proto/ballista.proto @@ -473,6 +473,9 @@ message FailedTask { // A successful task's result is lost due to executor lost ResultLost result_lost = 8; TaskKilled task_killed = 9; + // The task ran out of memory. Retriable: the retry may land on a less-loaded + // executor, or run once its peers have drained. + ResourcesExhausted resources_exhausted = 10; } } @@ -504,6 +507,9 @@ message ResultLost { message TaskKilled { } +message ResourcesExhausted { +} + message ShuffleWritePartition { uint64 partition_id = 1; uint64 num_batches = 3; diff --git a/ballista/core/src/error.rs b/ballista/core/src/error.rs index b0ce594f37..19c4e37fc0 100644 --- a/ballista/core/src/error.rs +++ b/ballista/core/src/error.rs @@ -24,7 +24,9 @@ use std::{ }; use crate::serde::protobuf::failed_task::FailedReason; -use crate::serde::protobuf::{ExecutionError, FailedTask, FetchPartitionError, IoError}; +use crate::serde::protobuf::{ + ExecutionError, FailedTask, FetchPartitionError, IoError, ResourcesExhausted, +}; use datafusion::error::DataFusionError; use datafusion::{arrow::error::ArrowError, sql::sqlparser::parser}; use futures::future::Aborted; @@ -245,6 +247,25 @@ impl From for FailedTask { failed_reason: Some(FailedReason::IoError(IoError {})), } } + BallistaError::DataFusionError(ref e) + if matches!(e.find_root(), DataFusionError::ResourcesExhausted(_)) => + { + FailedTask { + error: format!( + "Task failed due to exhausted resources (memory or spill capacity): {e}" + ), + // Retriable: the retry may land on a less-loaded executor, or run once + // this executor's other tasks have drained. Bounded by + // --task-max-failures, so a task that is simply too large still fails + // the job -- but with a clear message, rather than an OOM-killed + // executor and a FetchPartitionError cascade. + retryable: true, + count_to_failures: true, + failed_reason: Some(FailedReason::ResourcesExhausted( + ResourcesExhausted {}, + )), + } + } other => FailedTask { error: format!("Task failed due to runtime execution error: {other:?}"), retryable: false, @@ -256,3 +277,143 @@ impl From for FailedTask { } impl Error for BallistaError {} + +#[cfg(test)] +mod tests { + use super::*; + use crate::serde::protobuf::failed_task::FailedReason; + + #[test] + fn resources_exhausted_maps_to_a_retriable_failed_task() { + let err = BallistaError::DataFusionError(Box::new( + DataFusionError::ResourcesExhausted("over budget".to_string()), + )); + let failed: FailedTask = err.into(); + + assert!(failed.retryable, "an OOM'd task must be retried"); + assert!(failed.count_to_failures, "retries must be bounded"); + assert!( + matches!( + failed.failed_reason, + Some(FailedReason::ResourcesExhausted(_)) + ), + "expected ResourcesExhausted, got {:?}", + failed.failed_reason + ); + } + + #[test] + fn wrapped_resources_exhausted_is_still_recognized() { + // DataFusion routinely wraps errors in `Context`, so the mapping must look at + // the root cause rather than the outermost error. + let inner = DataFusionError::ResourcesExhausted("over budget".to_string()); + let err = BallistaError::DataFusionError(Box::new( + inner.context("while executing HashJoinExec"), + )); + let failed: FailedTask = err.into(); + + assert!(failed.retryable); + assert!(matches!( + failed.failed_reason, + Some(FailedReason::ResourcesExhausted(_)) + )); + } + + /// `DataFusionError::Shared` is how DataFusion fans *one* stream error out to + /// *many* output partitions: `RepartitionExec` and `CoalescePartitionsExec` clone a + /// single error into an `Arc` and hand it to every consumer. A memory rejection + /// raised inside a join below a `RepartitionExec` therefore reaches the shuffle + /// writer wrapped in exactly this variant -- so the whole feature's retriability + /// depends on `find_root()` seeing through it. It does (`Shared`'s `Error::source` + /// yields the inner `DataFusionError`), and this is what holds that true. + #[test] + fn a_shared_resources_exhausted_is_still_recognized() { + let err = BallistaError::DataFusionError(Box::new(DataFusionError::Shared( + std::sync::Arc::new(DataFusionError::ResourcesExhausted( + "over budget".to_string(), + )), + ))); + let failed: FailedTask = err.into(); + + assert!( + failed.retryable, + "a ResourcesExhausted fanned out through RepartitionExec must still be retried" + ); + assert!(failed.count_to_failures, "retries must be bounded"); + assert!( + matches!( + failed.failed_reason, + Some(FailedReason::ResourcesExhausted(_)) + ), + "expected ResourcesExhausted, got {:?}", + failed.failed_reason + ); + } + + /// The realistic shape: DataFusion adds a `Context` to the pool's rejection as it + /// unwinds out of the operator, and *then* the repartition shares it out. Both + /// wrappers have to be seen through. + #[test] + fn a_shared_context_wrapped_resources_exhausted_is_still_recognized() { + let inner = DataFusionError::ResourcesExhausted("over budget".to_string()); + let err = BallistaError::DataFusionError(Box::new(DataFusionError::Shared( + std::sync::Arc::new(inner.context("while executing HashJoinExec")), + ))); + let failed: FailedTask = err.into(); + + assert!(failed.retryable); + assert!(matches!( + failed.failed_reason, + Some(FailedReason::ResourcesExhausted(_)) + )); + } + + /// The discriminator for the two tests above: `Shared` must not become a blanket + /// "retriable" arm. A shared error whose root is *not* `ResourcesExhausted` still + /// has to fall through to the non-retriable execution-error arm. + #[test] + fn a_shared_non_resource_error_remains_non_retriable() { + let err = BallistaError::DataFusionError(Box::new(DataFusionError::Shared( + std::sync::Arc::new(DataFusionError::Execution("boom".to_string())), + ))); + let failed: FailedTask = err.into(); + + assert!(!failed.retryable); + assert!(matches!( + failed.failed_reason, + Some(FailedReason::ExecutionError(_)) + )); + } + + #[test] + fn other_errors_remain_non_retriable_execution_errors() { + let err = BallistaError::General("boom".to_string()); + let failed: FailedTask = err.into(); + + assert!(!failed.retryable); + assert!(matches!( + failed.failed_reason, + Some(FailedReason::ExecutionError(_)) + )); + } + + #[test] + fn other_datafusion_errors_remain_non_retriable_execution_errors() { + // Sharper discriminator than a bare `General` error: this is a + // `DataFusionError` whose root is NOT `ResourcesExhausted`, so it must + // still fall through to the non-retriable arm. It also happens to be + // the exact shape the old `format!("{e:?}")` re-wrap at + // shuffle_writer.rs used to produce, so it guards against the + // resources-exhausted arm accidentally matching every DataFusionError. + let err = BallistaError::DataFusionError(Box::new(DataFusionError::Execution( + "boom".to_string(), + ))); + let failed: FailedTask = err.into(); + + assert!(!failed.retryable); + assert!(matches!( + failed.failed_reason, + Some(FailedReason::ExecutionError(_)) + )); + } +} diff --git a/ballista/core/src/execution_plans/shuffle_writer.rs b/ballista/core/src/execution_plans/shuffle_writer.rs index 4a99197061..93aa8f7d4d 100644 --- a/ballista/core/src/execution_plans/shuffle_writer.rs +++ b/ballista/core/src/execution_plans/shuffle_writer.rs @@ -35,6 +35,7 @@ use std::sync::Arc; use std::time::Instant; use crate::JobId; +use crate::error::BallistaError; use crate::execution_plans::create_shuffle_path; use crate::extension::SessionConfigExt; use crate::utils; @@ -242,7 +243,15 @@ impl ShuffleWriterExec { channel_capacity, ) .await - .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?; + .map_err(|e| match e { + // Preserve the DataFusion error type: `find_root()` in the + // FailedTask mapping relies on it to classify a + // ResourcesExhausted as a retriable failure. A `format!` here + // would flatten it to an Execution error and silently make an + // OOM'd task non-retryable. + BallistaError::DataFusionError(e) => *e, + other => DataFusionError::Execution(format!("{other:?}")), + })?; write_metrics .input_rows @@ -600,10 +609,15 @@ fn result_schema() -> SchemaRef { #[allow(dead_code, unused_imports)] // clippy false positive with local imports mod tests { use super::*; + use crate::error::BallistaError; + use crate::serde::protobuf::FailedTask; + use crate::serde::protobuf::failed_task::FailedReason; use datafusion::arrow::array::{StringArray, StructArray, UInt32Array, UInt64Array}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; + use datafusion::physical_expr::EquivalenceProperties; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; + use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::expressions::Column; use datafusion::prelude::SessionContext; use tempfile::TempDir; @@ -803,6 +817,138 @@ mod tests { Ok(()) } + /// Test-only plan whose single stream immediately yields a + /// `DataFusionError::ResourcesExhausted`, standing in for an OOM'd input + /// (e.g. a hash-join build side or a sort spilling past its budget) that + /// feeds a `None`-partitioned (unpartitioned) shuffle write stage. + #[derive(Debug)] + struct AlwaysResourcesExhaustedExec { + properties: Arc, + schema: SchemaRef, + } + + impl AlwaysResourcesExhaustedExec { + fn new(schema: SchemaRef) -> Self { + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::clone(&schema)), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { properties, schema } + } + } + + impl DisplayAs for AlwaysResourcesExhaustedExec { + fn fmt_as( + &self, + t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + match t { + DisplayFormatType::Default + | DisplayFormatType::Verbose + | DisplayFormatType::TreeRender => { + write!(f, "AlwaysResourcesExhaustedExec") + } + } + } + } + + impl ExecutionPlan for AlwaysResourcesExhaustedExec { + fn name(&self) -> &str { + "AlwaysResourcesExhaustedExec" + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + let schema = self.schema(); + let stream = futures::stream::once(async { + Err(DataFusionError::ResourcesExhausted( + "over budget".to_string(), + )) + }); + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } + + fn partition_statistics( + &self, + _partition: Option, + ) -> Result> { + Ok(Arc::new(Statistics::new_unknown(&self.schema))) + } + } + + #[tokio::test] + async fn test_no_repart_resources_exhausted_is_retryable() -> Result<()> { + // This pins the writer boundary, not just the FailedTask mapping in + // error.rs: a `ShuffleWriterExec` with `shuffle_output_partitioning: + // None` (the final stage, a broadcast-join build side, or a + // CoalescePartitionsExec/SortPreservingMergeExec input) must let a + // `ResourcesExhausted` from its input stream survive the disk-write + // path with its error type intact, so the scheduler retries the task + // instead of failing the whole job. + let session_ctx = SessionContext::new(); + let task_ctx = session_ctx.task_ctx(); + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::UInt32, true)])); + let input_plan: Arc = + Arc::new(AlwaysResourcesExhaustedExec::new(schema)); + let work_dir = TempDir::new()?; + let query_stage = ShuffleWriterExec::try_new( + JobId::new("jobOne"), + 1, + input_plan, + work_dir.path().to_str().unwrap().to_owned(), + None, + )?; + let mut stream = query_stage.execute(0, task_ctx)?; + let result = utils::collect_stream(&mut stream).await; + + let err: BallistaError = result.expect_err( + "AlwaysResourcesExhaustedExec's stream error must propagate as an error", + ); + let failed: FailedTask = err.into(); + + assert!( + failed.retryable, + "a ResourcesExhausted task on a None-partitioned shuffle write stage \ + must be retryable, got: {failed:?}" + ); + assert!( + matches!( + failed.failed_reason, + Some(FailedReason::ResourcesExhausted(_)) + ), + "expected ResourcesExhausted, got {:?}", + failed.failed_reason + ); + + Ok(()) + } + fn create_input_plan() -> Result> { let schema = Arc::new(Schema::new(vec![ Field::new("a", DataType::UInt32, true), diff --git a/ballista/core/src/serde/generated/ballista.rs b/ballista/core/src/serde/generated/ballista.rs index 07dfa8a42f..ca83fc8968 100644 --- a/ballista/core/src/serde/generated/ballista.rs +++ b/ballista/core/src/serde/generated/ballista.rs @@ -717,7 +717,7 @@ pub struct FailedTask { /// Whether this task failure should be counted to the maximum number of times the task is allowed to retry #[prost(bool, tag = "3")] pub count_to_failures: bool, - #[prost(oneof = "failed_task::FailedReason", tags = "4, 5, 6, 7, 8, 9")] + #[prost(oneof = "failed_task::FailedReason", tags = "4, 5, 6, 7, 8, 9, 10")] pub failed_reason: ::core::option::Option, } /// Nested message and enum types in `FailedTask`. @@ -737,6 +737,10 @@ pub mod failed_task { ResultLost(super::ResultLost), #[prost(message, tag = "9")] TaskKilled(super::TaskKilled), + /// The task ran out of memory. Retriable: the retry may land on a less-loaded + /// executor, or run once its peers have drained. + #[prost(message, tag = "10")] + ResourcesExhausted(super::ResourcesExhausted), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -768,6 +772,8 @@ pub struct ResultLost {} #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] pub struct TaskKilled {} #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] +pub struct ResourcesExhausted {} +#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] pub struct ShuffleWritePartition { #[prost(uint64, tag = "1")] pub partition_id: u64, diff --git a/ballista/scheduler/src/api/handlers.rs b/ballista/scheduler/src/api/handlers.rs index 8da872253a..f1ab16c1da 100644 --- a/ballista/scheduler/src/api/handlers.rs +++ b/ballista/scheduler/src/api/handlers.rs @@ -23,7 +23,8 @@ use axum::{ response::{IntoResponse, Response}, }; use ballista_core::serde::protobuf::failed_task::FailedReason::{ - ExecutionError, ExecutorLost, FetchPartitionError, IoError, ResultLost, TaskKilled, + ExecutionError, ExecutorLost, FetchPartitionError, IoError, ResourcesExhausted, + ResultLost, TaskKilled, }; use ballista_core::serde::protobuf::job_status::Status; use ballista_core::serde::protobuf::{ @@ -823,6 +824,7 @@ fn failed_reason(failed: &FailedTask) -> String { Some(ExecutorLost(_)) => "ExecutorLost", Some(ResultLost(_)) => "ResultLost", Some(TaskKilled(_)) => "TaskKilled", + Some(ResourcesExhausted(_)) => "ResourcesExhausted", None => "Failed", } .to_string() From bbf8a6f6a292458ce794ce130b1411283ed6f523 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 13 Jul 2026 14:23:43 -0600 Subject: [PATCH 2/2] feat(executor): add an allocator-backed OOM guard behind the oom-guard feature Executor memory accounting relies on voluntary MemoryPool reservations, which miss allocations made by Arrow buffers, join scratch space, and expression kernels. Real native memory therefore runs above what the pool believes is reserved, and the executor process can be OOM-killed. That is far worse than a single failed task: a dead executor takes all of its shuffle output with it, so every downstream stage that would have read from it raises a fetch-partition error, cascading into stage rollbacks and map-stage re-runs, potentially across other concurrent jobs. Add a global allocator that tracks the bytes actually handed out, and use that signal in two layers: - RealUsagePool, a MemoryPool decorator that rejects growth once real live usage plus the request would exceed the budget, so DataFusion spills rather than growing into an OOM. This is the layer that does the work. - MemoryGuardExec, a transparent pass-through plan node on each stage's data path that checks the budget between batches and fails that one task as a last resort. It debounces, so a transient spike that the pool gate is already resolving does not shed every task on the executor. The allocator itself only tracks: unwinding out of a GlobalAlloc is undefined behaviour, so enforcement happens at safe points instead. Live bytes are counted in cache-line-padded shards, which keeps the count exact across thread churn without a Drop-carrying thread-local (whose lazy storage would allocate from inside the allocator). The guard enforces on tracked live bytes rather than RSS. RSS is what the OOM killer reads, but an allocator does not promptly return freed pages, so a spill that genuinely released memory would leave RSS high and the guard latched on, failing every subsequent task. Failing tasks cannot reclaim memory that is already free. RSS is sampled for observability only, and a warning is logged when it diverges above the limit while tracked usage stays below it, which is the signal that the allocator is holding memory the guard cannot reclaim. The feature is off by default. It installs a tracking global allocator, so it is a build-time opt-in and the default build is unchanged, with no per-allocation overhead. --- ballista/executor/Cargo.toml | 3 + ballista/executor/src/bin/main.rs | 18 +- ballista/executor/src/config.rs | 29 +- ballista/executor/src/execution_engine.rs | 354 ++++- ballista/executor/src/executor_process.rs | 43 + ballista/executor/src/lib.rs | 3 + .../executor/src/memory_pools/guard_exec.rs | 576 ++++++++ ballista/executor/src/memory_pools/mod.rs | 49 + .../executor/src/memory_pools/oom_guard.rs | 1154 +++++++++++++++++ .../src/memory_pools/real_usage_pool.rs | 316 +++++ ballista/executor/tests/oom_guard_alloc.rs | 284 ++++ .../user-guide/deployment/cargo-install.md | 15 + docs/source/user-guide/tuning-guide.md | 74 ++ 13 files changed, 2899 insertions(+), 19 deletions(-) create mode 100644 ballista/executor/src/memory_pools/guard_exec.rs create mode 100644 ballista/executor/src/memory_pools/mod.rs create mode 100644 ballista/executor/src/memory_pools/oom_guard.rs create mode 100644 ballista/executor/src/memory_pools/real_usage_pool.rs create mode 100644 ballista/executor/tests/oom_guard_alloc.rs diff --git a/ballista/executor/Cargo.toml b/ballista/executor/Cargo.toml index caae3b4260..4170c73fff 100644 --- a/ballista/executor/Cargo.toml +++ b/ballista/executor/Cargo.toml @@ -36,6 +36,9 @@ required-features = ["build-binary"] arrow-ipc-optimizations = [] build-binary = ["clap", "tracing-subscriber", "tracing-appender", "tracing", "ballista-core/build-binary", "mimalloc"] default = ["arrow-ipc-optimizations", "build-binary"] +# Allocator-backed OOM protection. Installs a tracking global allocator, so it is +# opt-in and off by default; the default build has zero per-allocation overhead. +oom-guard = [] spark-compat = ["ballista-core/spark-compat"] [dependencies] diff --git a/ballista/executor/src/bin/main.rs b/ballista/executor/src/bin/main.rs index cbeb93b89e..8a7912fa9e 100644 --- a/ballista/executor/src/bin/main.rs +++ b/ballista/executor/src/bin/main.rs @@ -30,10 +30,26 @@ use std::env; use std::sync::Arc; use tracing_subscriber::EnvFilter; -#[cfg(feature = "mimalloc")] +#[cfg(all(feature = "mimalloc", not(feature = "oom-guard")))] #[global_allocator] static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc; +#[cfg(all(feature = "oom-guard", feature = "mimalloc"))] +#[global_allocator] +static GLOBAL: ballista_executor::memory_pools::oom_guard::AccountingAllocator< + mimalloc::MiMalloc, +> = ballista_executor::memory_pools::oom_guard::AccountingAllocator::new( + mimalloc::MiMalloc, +); + +#[cfg(all(feature = "oom-guard", not(feature = "mimalloc")))] +#[global_allocator] +static GLOBAL: ballista_executor::memory_pools::oom_guard::AccountingAllocator< + std::alloc::System, +> = ballista_executor::memory_pools::oom_guard::AccountingAllocator::new( + std::alloc::System, +); + #[tokio::main] async fn main() -> ballista_core::error::Result<()> { // parse command-line arguments diff --git a/ballista/executor/src/config.rs b/ballista/executor/src/config.rs index 48b5a28c9c..bdb0c5afa3 100644 --- a/ballista/executor/src/config.rs +++ b/ballista/executor/src/config.rs @@ -165,10 +165,31 @@ pub struct Config { /// Optional total memory budget for the executor. Accepts human-readable /// values like "8GB", "512MiB", or a plain byte count. When set, every /// task gets a FairSpillPool of size `memory_pool_size / concurrent_tasks`. - #[arg( - long, - value_parser = parse_memory_pool_size, - help = "Optional total executor memory budget (e.g. \"8GB\", \"512MiB\"). Each concurrent task receives an equal share." + /// + /// With the `oom-guard` feature compiled in, this value is *also* the ceiling + /// the tracking allocator enforces, and that ceiling covers every live byte the + /// process has allocated -- gRPC buffers, tokio, object-store caches, shuffle + /// writes -- not just query memory. Expect an `oom-guard` build to start + /// spilling somewhat earlier at a given budget than the default build does. + // + // The help text is `cfg`-selected: a default-build executor has no tracking + // allocator, so telling its operator about a feature that is not compiled in would + // simply be wrong. + #[cfg_attr( + not(feature = "oom-guard"), + arg( + long, + value_parser = parse_memory_pool_size, + help = "Optional total executor memory budget (e.g. \"8GB\", \"512MiB\"). Each concurrent task receives an equal share." + ) + )] + #[cfg_attr( + feature = "oom-guard", + arg( + long, + value_parser = parse_memory_pool_size, + help = "Optional total executor memory budget (e.g. \"8GB\", \"512MiB\"). Each concurrent task receives an equal share. This build has the oom-guard feature, so the value is also the process-wide ceiling on live allocator bytes (all process memory, not just query memory) and therefore gates somewhat earlier at the same value." + ) )] pub memory_pool_size: Option, /// Maximum number of sessions whose shared runtime state (object-store diff --git a/ballista/executor/src/execution_engine.rs b/ballista/executor/src/execution_engine.rs index e579621c9e..81495f9eb8 100644 --- a/ballista/executor/src/execution_engine.rs +++ b/ballista/executor/src/execution_engine.rs @@ -165,16 +165,71 @@ fn restrict_scan_to_partition( Some(DataSourceExec::from_data_source(config)) } -impl ExecutionEngine for DefaultExecutionEngine { - fn create_query_stage_exec( +/// The plan the stage's shuffle writer consumes: the writer's own child, wrapped in a +/// [`MemoryGuardExec`](crate::memory_pools::MemoryGuardExec). +/// +/// The guard's stream checks the executor's real allocator balance against the armed +/// limit before every batch, and fails this one task with a retriable +/// `ResourcesExhausted` rather than letting the process be OOM-killed. Placing it here +/// puts the check on the stage's entire data path. +/// +/// **The insertion point is load-bearing**, for two reasons -- neither of which is that +/// the guard hides the plan beneath it. It does not: `MemoryGuardExec` implements no +/// `downcast_delegate`, so `downcast_ref` on the guard itself yields `None`, but +/// `TreeNode::transform` descends through `children()` / `with_new_children()` and +/// downcasts each node it visits individually. An opaque node conceals nothing below +/// itself from a traversal. (DataFusion says as much of `downcast_delegate`: it "should +/// not be used for plan traversal or optimizer rewrites".) +/// +/// What actually matters: +/// +/// 1. **The guard must not sit on the stage root.** [`build_stage_writer`] and the +/// scheduler both `downcast_ref` the root to find the shuffle writer +/// (`ShuffleWriterExec` / `SortShuffleWriterExec`). A guard *there* is +/// opaque to that downcast and the stage would not be recognised as a writer at all. +/// Hence the guard goes immediately *below* the writer, wrapping the writer's child. +/// See `the_guard_is_inserted_immediately_below_each_shuffle_writer`. +/// 2. **Exactly one guard must be inserted.** The guard adds a level to the plan, and +/// the executor's flattened metrics list must stay in step with the scheduler's view +/// of the same plan, which is zipped by position. One guard, at a known depth, keeps +/// that predictable. See `exactly_one_guard_is_inserted` and +/// `the_guard_does_not_change_the_metrics_list_length`. +/// +/// Do not give `MemoryGuardExec` a `downcast_delegate`, and do not move this call onto +/// the root. +/// +/// [`build_stage_writer`]: DefaultExecutionEngine::build_stage_writer +#[cfg(feature = "oom-guard")] +fn stage_input(plan: &Arc) -> Arc { + Arc::new(crate::memory_pools::MemoryGuardExec::new( + plan.children()[0].clone(), + )) +} + +/// The plan the stage's shuffle writer consumes: the writer's own child, unchanged. +/// +/// Without the `oom-guard` feature there is no tracking allocator and nothing to check, +/// so the stage plan is rebuilt exactly as before. +#[cfg(not(feature = "oom-guard"))] +fn stage_input(plan: &Arc) -> Arc { + plan.children()[0].clone() +} + +impl DefaultExecutionEngine { + /// Rebuild the scheduler's stage plan for execution on this executor and return its + /// shuffle writer. + /// + /// Split out of [`ExecutionEngine::create_query_stage_exec`] so that tests can assert + /// on the concrete rewritten plan (which the `Arc` it is + /// wrapped in hides). + fn build_stage_writer( &self, job_id: JobId, stage_id: usize, partition_id: usize, plan: Arc, work_dir: &str, - _config: &SessionConfig, - ) -> Result> { + ) -> Result { let plan = plan .transform(|p| { if let Some(reader) = p.downcast_ref::() { @@ -205,13 +260,11 @@ impl ExecutionEngine for DefaultExecutionEngine { let exec = ShuffleWriterExec::try_new( job_id, stage_id, - plan.children()[0].clone(), + stage_input(&plan), work_dir.to_string(), shuffle_writer.shuffle_output_partitioning().cloned(), )?; - Ok(Arc::new(DefaultQueryStageExec::new( - ShuffleWriterVariant::Hash(exec), - ))) + Ok(ShuffleWriterVariant::Hash(exec)) } else if let Some(sort_shuffle_writer) = plan.downcast_ref::() { @@ -219,14 +272,12 @@ impl ExecutionEngine for DefaultExecutionEngine { let exec = SortShuffleWriterExec::try_new( job_id, stage_id, - plan.children()[0].clone(), + stage_input(&plan), work_dir.to_string(), sort_shuffle_writer.shuffle_output_partitioning().clone(), sort_shuffle_writer.config().clone(), )?; - Ok(Arc::new(DefaultQueryStageExec::new( - ShuffleWriterVariant::Sort(exec), - ))) + Ok(ShuffleWriterVariant::Sort(exec)) } else { Err(DataFusionError::Internal( "Plan passed to new_query_stage_exec is not a ShuffleWriterExec or SortShuffleWriterExec" @@ -236,6 +287,22 @@ impl ExecutionEngine for DefaultExecutionEngine { } } +impl ExecutionEngine for DefaultExecutionEngine { + fn create_query_stage_exec( + &self, + job_id: JobId, + stage_id: usize, + partition_id: usize, + plan: Arc, + work_dir: &str, + _config: &SessionConfig, + ) -> Result> { + let writer = + self.build_stage_writer(job_id, stage_id, partition_id, plan, work_dir)?; + Ok(Arc::new(DefaultQueryStageExec::new(writer))) + } +} + /// Enum representing the different shuffle writer implementations. #[derive(Debug, Clone)] pub enum ShuffleWriterVariant { @@ -331,15 +398,104 @@ impl QueryStageExecutor for DefaultQueryStageExec { #[cfg(test)] mod tests { use super::*; - use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use ballista_core::client_pool::{BallistaClientPool, PooledClient}; + #[cfg(feature = "oom-guard")] + use ballista_core::execution_plans::sort_shuffle::SortShuffleConfig; + use ballista_core::extension::BallistaConfigGrpcEndpoint; + use ballista_core::utils::GrpcClientConfig; + use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::datasource::listing::PartitionedFile; use datafusion::datasource::physical_plan::ParquetSource; use datafusion::execution::object_store::ObjectStoreUrl; + #[cfg(feature = "oom-guard")] + use datafusion::physical_expr::expressions::col; + use datafusion::physical_plan::Partitioning; use datafusion::physical_plan::empty::EmptyExec; + use datafusion::physical_plan::union::UnionExec; + + /// The schema every plan in these tests is built over. + fn test_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])) + } + + /// The work dir the rewritten stage plan must be given. + const WORK_DIR: &str = "/work/dir"; + + /// A stage plan as the scheduler hands it over: rooted at a `ShuffleWriterExec`. + fn hash_writer(input: Arc) -> Arc { + Arc::new( + ShuffleWriterExec::try_new( + "job".into(), + 1, + input, + "/scheduler/dir".to_string(), + None, + ) + .unwrap(), + ) + } + + /// The same, rooted at a `SortShuffleWriterExec` -- the branch that is easy to forget. + /// Only the guard-placement tests build this variant, and they only exist when the + /// guard is compiled in. + #[cfg(feature = "oom-guard")] + fn sort_writer(input: Arc) -> Arc { + let partitioning = Partitioning::Hash(vec![col("a", &test_schema()).unwrap()], 2); + Arc::new( + SortShuffleWriterExec::try_new( + "job".into(), + 1, + input, + "/scheduler/dir".to_string(), + partitioning, + SortShuffleConfig::default(), + ) + .unwrap(), + ) + } + + /// The rebuilt shuffle writer, as an `ExecutionPlan`, whichever variant it is. + fn writer_root(writer: &ShuffleWriterVariant) -> Arc { + match writer { + ShuffleWriterVariant::Hash(w) => Arc::new(w.clone()), + ShuffleWriterVariant::Sort(w) => Arc::new(w.clone()), + } + } + + /// The first node of type `T` anywhere in the tree, in pre-order. + fn find_first( + plan: &Arc, + ) -> Option> { + if plan.downcast_ref::().is_some() { + return Some(Arc::clone(plan)); + } + plan.children().into_iter().find_map(find_first::) + } + + /// A client pool test double with a recognizable `Debug` rendering: the reader's + /// `client_pool` field is private, so `Debug` is how a test observes that the + /// `transform` pass installed it. + #[derive(Debug)] + struct TestClientPool; + + #[async_trait::async_trait] + impl BallistaClientPool for TestClientPool { + async fn acquire( + &self, + _host: &str, + _port: u16, + _config: &GrpcClientConfig, + _customize_endpoint: Option>, + ) -> ballista_core::error::Result { + unimplemented!("no client is ever acquired in these tests") + } + + async fn evict_idle(&self) {} + } /// Build a `DataSourceExec` over `n` file groups, one file each. fn scan_with_file_groups(n: usize) -> Arc { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let schema = test_schema(); let source = Arc::new(ParquetSource::new(schema)); let mut builder = FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source); @@ -383,4 +539,174 @@ mod tests { let plan: Arc = Arc::new(EmptyExec::new(schema)); assert!(restrict_scan_to_partition(&plan, 0).is_none()); } + + /// Inserting the guard must not disturb the `transform` pass: the + /// `ShuffleReaderExec` still receives its `work_dir` / `client_pool` (without which + /// the stage cannot read its shuffle input) and the `DataSourceExec` is still + /// restricted to this task's partition (without which the task over-reads the shared + /// file group -- silently wrong results, not an error). Runs in both feature + /// configurations. + /// + /// Note what this does *not* prove. A guard hoisted above the `transform` pass would + /// still leave both rewrites intact, because `TreeNode::transform` walks the plan via + /// `children()` / `with_new_children()` -- which `MemoryGuardExec` implements -- and + /// downcasts each node individually, so an opaque node in the path hides nothing + /// beneath it. The guard's opacity to `downcast_ref` only breaks a downcast aimed at a + /// node the guard sits *directly on top of* -- in this function, the shuffle writer on + /// the stage root. That is what `the_guard_is_inserted_immediately_below_each_shuffle_writer` + /// and `exactly_one_guard_is_inserted` pin down; both fail if the guard is inserted + /// anywhere but immediately below the writer. + #[test] + fn the_transform_pass_still_reaches_the_reader_and_the_scan() { + let engine = DefaultExecutionEngine::with_client_pool(Arc::new(TestClientPool)); + + let reader: Arc = Arc::new( + ShuffleReaderExec::try_new( + 1, + vec![vec![]], + test_schema(), + Partitioning::UnknownPartitioning(1), + ) + .unwrap(), + ); + let union: Arc = + UnionExec::try_new(vec![reader, scan_with_file_groups(4)]).unwrap(); + + let writer = engine + .build_stage_writer("job".into(), 2, 2, hash_writer(union), WORK_DIR) + .unwrap(); + let root = writer_root(&writer); + + let reader = find_first::(&root) + .expect("the rebuilt plan must still contain the shuffle reader"); + // `work_dir` and `client_pool` are private to `ShuffleReaderExec`; its derived + // `Debug` is the only view a test outside `ballista-core` has of them. + let rendered = format!("{reader:?}"); + assert!( + rendered.contains(&format!("work_dir: Some({WORK_DIR:?})")), + "the reader must have been given its work_dir, got: {rendered}" + ); + assert!( + rendered.contains("client_pool: Some(TestClientPool)"), + "the reader must have been given the client pool, got: {rendered}" + ); + + let scan = find_first::(&root) + .expect("the rebuilt plan must still contain the scan"); + assert_eq!( + group_file_counts(&scan), + vec![0, 0, 1, 0], + "the scan must be restricted to this task's partition (2 of 4); an \ + unrestricted scan over-reads the shared file group -- wrong results" + ); + } + + /// `collect_plan_metrics` pushes a node's metrics *and* recurses into its children, + /// and the scheduler's `merge_stage_metrics` positionally zips the executor's list + /// against its own guard-free plan -- silently dropping the metrics of every stage of + /// every job on a length mismatch. `MemoryGuardExec` therefore reports no metrics of + /// its own, and the two lists must stay the same length. This fails the instant + /// someone adds a `metrics()` delegation to the guard. + #[test] + fn the_guard_does_not_change_the_metrics_list_length() { + let engine = DefaultExecutionEngine::new(); + // The plan as the scheduler holds it: no guard, ever. + let scheduler_plan = hash_writer(scan_with_file_groups(2)); + let writer = engine + .build_stage_writer("job".into(), 1, 0, Arc::clone(&scheduler_plan), WORK_DIR) + .unwrap(); + let executor_plan = writer_root(&writer); + + assert_eq!( + utils::collect_plan_metrics(executor_plan.as_ref()).len(), + utils::collect_plan_metrics(scheduler_plan.as_ref()).len(), + "the executor's flattened metrics list must stay the same length as the \ + scheduler's guard-free view of the same plan" + ); + } + + /// Number of `MemoryGuardExec` nodes anywhere in the tree. + #[cfg(feature = "oom-guard")] + fn count_guards(plan: &Arc) -> usize { + use crate::memory_pools::MemoryGuardExec; + usize::from(plan.downcast_ref::().is_some()) + + plan.children().into_iter().map(count_guards).sum::() + } + + /// The guard sits immediately below the shuffle writer -- and nowhere else. The root + /// must stay an un-wrapped writer (the scheduler's own downcast, and this function's, + /// depend on it), and the guard's child must be the original input. + /// + /// Asserted for *both* writer variants: the `SortShuffleWriterExec` branch is easy to + /// forget. + #[cfg(feature = "oom-guard")] + #[test] + fn the_guard_is_inserted_immediately_below_each_shuffle_writer() { + use crate::memory_pools::MemoryGuardExec; + + let engine = DefaultExecutionEngine::new(); + let input = Arc::new(EmptyExec::new(test_schema())) as Arc; + + for (label, plan) in [ + ("ShuffleWriterExec", hash_writer(Arc::clone(&input))), + ("SortShuffleWriterExec", sort_writer(Arc::clone(&input))), + ] { + let writer = engine + .build_stage_writer("job".into(), 1, 0, plan, WORK_DIR) + .unwrap(); + let root = writer_root(&writer); + + assert!( + root.downcast_ref::().is_none(), + "{label}: the stage root must remain the writer, not a guard -- \ + `create_query_stage_exec` downcasts the root to find the writer" + ); + + let children = root.children(); + assert_eq!(children.len(), 1, "{label}: the writer has one child"); + let guard = children[0] + .downcast_ref::() + .unwrap_or_else(|| { + panic!("{label}: the writer's child must be a MemoryGuardExec") + }); + assert!( + Arc::ptr_eq(guard.input(), &input), + "{label}: the guard must wrap the writer's original child, unchanged" + ); + } + } + + /// Exactly one guard, ever: no guard-over-guard, and none deeper in the tree. Every + /// downcast site assumes the guard lives at exactly one known depth. + #[cfg(feature = "oom-guard")] + #[test] + fn exactly_one_guard_is_inserted() { + let engine = DefaultExecutionEngine::new(); + let reader: Arc = Arc::new( + ShuffleReaderExec::try_new( + 1, + vec![vec![]], + test_schema(), + Partitioning::UnknownPartitioning(1), + ) + .unwrap(), + ); + let union: Arc = + UnionExec::try_new(vec![reader, scan_with_file_groups(4)]).unwrap(); + + for plan in [ + hash_writer(Arc::clone(&union)), + sort_writer(Arc::clone(&union)), + ] { + let writer = engine + .build_stage_writer("job".into(), 1, 0, plan, WORK_DIR) + .unwrap(); + let root = writer_root(&writer); + assert_eq!( + count_guards(&root), + 1, + "exactly one guard must be inserted, immediately below the writer" + ); + } + } } diff --git a/ballista/executor/src/executor_process.rs b/ballista/executor/src/executor_process.rs index 360a894913..fdcecfe9f7 100644 --- a/ballista/executor/src/executor_process.rs +++ b/ballista/executor/src/executor_process.rs @@ -99,6 +99,13 @@ fn memory_pool_policy( Ok(Arc::new( move |base: Arc, _config: &SessionConfig| { let pool: Arc = Arc::new(FairSpillPool::new(per_task)); + // The tracked-reservation pool above only sees explicit reservations. Wrap it + // so growth is also gated on the *real* allocator balance against the whole + // executor budget, which catches untracked Arrow / join / kernel bytes. + #[cfg(feature = "oom-guard")] + let pool: Arc = Arc::new( + crate::memory_pools::RealUsagePool::new(pool, total_bytes as usize), + ); RuntimeEnvBuilder::from_runtime_env(&base) .with_memory_pool(pool) .build_arc() @@ -322,8 +329,44 @@ pub async fn start_executor_process( info!( "Memory pool: total {total} bytes split into {concurrent_tasks} tasks ({per_task} bytes each)" ); + // Set the breaker's limit to the same threshold as the cooperative gate. The two + // order correctly: the gate trips on *projected* usage (balance + request) so it + // spills first; the breaker trips on *actual* usage as the backstop. + #[cfg(feature = "oom-guard")] + { + crate::memory_pools::oom_guard::arm(total as usize); + info!("OOM guard armed at {total} bytes of tracked live allocator usage"); + + // Advisory only: this ticker samples RSS and *logs*, it never writes the + // enforced balance. Enforcement reads the allocator's exact live-bytes + // counter, which drops the instant memory is freed -- so a spill un-trips the + // guard immediately. Feeding RSS back into that counter would be actively + // harmful: an allocator like mimalloc `MADV_FREE`s pages but leaves them + // resident, so RSS stays high after a successful spill and would re-raise the + // balance over the limit, latching the guard into failing every batch of + // every task on this executor. What RSS is good for is telling an operator + // *which* of those two situations they are in, which is what this logs. + tokio::spawn(async move { + let mut ticker = time::interval(Duration::from_secs(1)); + loop { + ticker.tick().await; + crate::memory_pools::oom_guard::observe_rss( + memory_stats::memory_stats().map(|usage| usage.physical_mem), + ); + } + }); + } policy } else { + // Two log lines, `cfg`-selected: the default build has no guard at all, and + // mentioning one there would send an operator hunting for something that does + // not exist in their binary. + #[cfg(feature = "oom-guard")] + info!( + "Memory pool: unbounded (--memory-pool-size not set); the OOM guard tracks allocations but does not enforce a limit" + ); + #[cfg(not(feature = "oom-guard"))] + info!("Memory pool: unbounded (--memory-pool-size not set)"); identity_pool_policy() }; diff --git a/ballista/executor/src/lib.rs b/ballista/executor/src/lib.rs index 70234ceaa5..37593de9de 100644 --- a/ballista/executor/src/lib.rs +++ b/ballista/executor/src/lib.rs @@ -37,6 +37,9 @@ pub mod executor_process; pub mod executor_server; /// Arrow Flight service for streaming shuffle data between executors. pub mod flight_service; +/// Allocator-backed OOM protection: real-usage memory accounting and circuit breaker. +#[cfg(feature = "oom-guard")] +pub mod memory_pools; /// Metrics collection for executor runtime statistics. pub mod metrics; /// Session-scoped cache of shared executor runtime environments. diff --git a/ballista/executor/src/memory_pools/guard_exec.rs b/ballista/executor/src/memory_pools/guard_exec.rs new file mode 100644 index 0000000000..301e3a0508 --- /dev/null +++ b/ballista/executor/src/memory_pools/guard_exec.rs @@ -0,0 +1,576 @@ +// 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. + +//! The circuit-breaker plan node. +//! +//! [`MemoryGuardExec`] is the last line of defence in the executor's OOM protection. +//! The cooperative gate in `RealUsagePool` makes DataFusion spill *before* the process +//! runs out of memory; this node handles the case where that is not enough -- memory +//! the pool never sees, growing past the limit anyway -- by failing the single offending +//! task with a retriable `ResourcesExhausted` instead of letting the OS OOM-kill the +//! whole executor and every task on it. +//! +//! Enforcement happens here, at a `poll_next` boundary, rather than in the allocator: +//! unwinding out of a `GlobalAlloc` is undefined behaviour, and DataFusion runs parts of +//! a plan on spawned tokio sub-tasks, where a panic would merely surface as a `JoinError` +//! and never reach a catch site. A poll boundary is a safe point, needs no unwinding +//! machinery, and costs a handful of relaxed atomic loads per batch. +//! +//! The quantity it checks is the allocator's *live-bytes* counter, which decrements on +//! every `dealloc`. So when the cooperative gate does its job and a consumer spills, the +//! guard stops tripping immediately: it can never latch into failing every batch of every +//! task on the executor. See `oom_guard` for why enforcing on RSS would do exactly that. + +use crate::memory_pools::oom_guard; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::{Result, Statistics, internal_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + SendableRecordBatchStream, +}; +use futures::StreamExt; +use std::fmt::Formatter; +use std::sync::Arc; + +/// A pass-through [`ExecutionPlan`] that fails its task when the executor's real memory +/// usage is over budget. +/// +/// The node is deliberately transparent: it forwards its input's schema, properties and +/// statistics untouched, and its stream yields the input's batches in order, unbuffered +/// and unmodified. Its only added behaviour is a call to [`oom_guard::check_budget`] +/// before each item the input yields; when that reports the process is over the limit, +/// the stream yields `ResourcesExhausted` instead of the batch, failing this one task. +/// +/// Metrics are the one exception to "forwards untouched": `metrics()` reports `None` +/// rather than delegating to the input, because Ballista's metrics collection already +/// visits the child separately (see the `metrics()` override below for why delegating +/// would double-count it). +/// +/// Transparency is a correctness requirement, not a nicety. Because +/// [`properties`](ExecutionPlan::properties) returns the input's `PlanProperties` +/// unchanged, inserting this node cannot alter output partitioning, ordering, emission +/// type or boundedness -- so it cannot change query results or perturb the scheduler's +/// stage planning. +#[derive(Debug)] +pub struct MemoryGuardExec { + /// The wrapped input; this node adds nothing to its output but the budget check. + input: Arc, +} + +impl MemoryGuardExec { + /// Wrap `input` so that every batch it produces is gated on the memory budget. + pub fn new(input: Arc) -> Self { + Self { input } + } + + /// The wrapped input plan. + pub fn input(&self) -> &Arc { + &self.input + } +} + +impl DisplayAs for MemoryGuardExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + match t { + DisplayFormatType::Default + | DisplayFormatType::Verbose + | DisplayFormatType::TreeRender => write!(f, "MemoryGuardExec"), + } + } +} + +impl ExecutionPlan for MemoryGuardExec { + fn name(&self) -> &str { + "MemoryGuardExec" + } + + fn schema(&self) -> SchemaRef { + self.input.schema() + } + + /// The input's properties, verbatim. Returning anything else would let this node + /// change the plan's partitioning or ordering -- a silent correctness bug. + fn properties(&self) -> &Arc { + self.input.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + if children.len() != 1 { + return internal_err!( + "MemoryGuardExec expected one child, got {}", + children.len() + ); + } + Ok(Arc::new(Self::new(children.pop().unwrap()))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let schema = self.input.schema(); + // Check before passing each item on. The check sums the balance shards with + // relaxed loads, so per batch is the right granularity: cheap enough to be free + // against the cost of producing a batch, frequent enough to stop a runaway task + // well before the process is killed -- and, because the balance falls as soon as + // memory is freed, to *stop* failing the moment a spill has done its job. + // + // Only gate the success path: if the child itself yielded an error, that error + // must win over an over-budget verdict. Using `Result::and` here would let a + // budget check that fails on the same item silently replace the child's real + // error (e.g. a corrupt-file or serde failure) with `ResourcesExhausted` -- + // misreporting a non-retriable failure as retriable and burning the task's + // retry budget on a doomed retry. + let guarded = self + .input + .execute(partition, context)? + .map(|item| item.and_then(|batch| oom_guard::check_budget().map(|_| batch))); + + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, guarded))) + } + + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + // Deliberately no `metrics()` override here: this node reports `None` (DataFusion's + // default), not the child's `MetricsSet`. `collect_plan_metrics` + // (`ballista/core/src/utils.rs`) pushes `plan.metrics()` for a node and then + // recurses into its children -- so delegating to `self.input.metrics()` would push + // the child's metrics twice (once for this node, once for the child itself). That + // desyncs the executor's flattened metrics list from the scheduler's metrics-free + // view of the same plan, which zips the two by position + // (`merge_stage_metrics` in `ballista/scheduler/src/display.rs`) -- a length + // mismatch makes it bail out, silently dropping stage metrics for every stage of + // every job. This node has no metrics of its own; the child's are already visited + // by the recursion. Do not re-add this delegation. + + fn partition_statistics(&self, partition: Option) -> Result> { + self.input.partition_statistics(partition) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::Equal + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::memory_pools::oom_guard; + use datafusion::arrow::array::Int32Array; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::common::DataFusionError; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::execution::TaskContext; + use datafusion::physical_expr::EquivalenceProperties; + use datafusion::physical_plan::Partitioning; + use datafusion::physical_plan::common::collect; + use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; + + fn test_input() -> Arc { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None).unwrap() + } + + /// The schema shared by the test doubles below: a single non-null `Int32` column. + fn test_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])) + } + + fn int_batch(schema: &SchemaRef, values: &[i32]) -> RecordBatch { + RecordBatch::try_new( + Arc::clone(schema), + vec![Arc::new(Int32Array::from(values.to_vec()))], + ) + .unwrap() + } + + /// A test-only input that streams three distinct batches on a single partition and, + /// as a side effect of producing the *second* one, drives the shared balance over + /// whatever limit is currently armed. It sets the balance directly (the production + /// path moves it only through the allocator's deltas) so the step is deterministic. + /// + /// This pins the "checks before every batch" contract from `MemoryGuardExec`: with + /// the guard's per-batch check, the first batch must still arrive intact (the budget + /// was fine when it was produced), and the stream must then fail once the second + /// batch's production has pushed the balance over the limit. A guard that instead + /// checked the budget once, eagerly, at `execute()` time -- before any batch has been + /// pulled from this input, and therefore before the side effect below ever runs -- + /// would let all three batches through, because the check would already be behind it + /// by the time the balance moves. That eager, one-shot behaviour is exactly what this + /// test must catch. + #[derive(Debug)] + struct StepBudgetExec { + properties: Arc, + schema: SchemaRef, + /// The balance value written once the second batch is produced. + raised_balance: usize, + } + + impl StepBudgetExec { + fn new(raised_balance: usize) -> Self { + let schema = test_schema(); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::clone(&schema)), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { + properties, + schema, + raised_balance, + } + } + } + + impl DisplayAs for StepBudgetExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + match t { + DisplayFormatType::Default + | DisplayFormatType::Verbose + | DisplayFormatType::TreeRender => write!(f, "StepBudgetExec"), + } + } + } + + impl ExecutionPlan for StepBudgetExec { + fn name(&self) -> &str { + "StepBudgetExec" + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if !children.is_empty() { + return internal_err!( + "StepBudgetExec expected no children, got {}", + children.len() + ); + } + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + let schema = self.schema(); + let batches = vec![ + int_batch(&schema, &[1, 2, 3]), + int_batch(&schema, &[4, 5, 6]), + int_batch(&schema, &[7, 8, 9]), + ]; + let raised_balance = self.raised_balance; + let stream = futures::stream::iter(batches.into_iter().enumerate()).map( + move |(i, batch)| { + if i == 1 { + // Fires while producing the *second* batch -- after the first + // has already been handed to the guard, never before `execute` + // was called. + oom_guard::test_support::set_balance_for_test( + raised_balance as isize, + ); + } + Ok(batch) + }, + ); + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } + + fn partition_statistics( + &self, + _partition: Option, + ) -> Result> { + Ok(Arc::new(Statistics::new_unknown(&self.schema))) + } + } + + /// A test-only input whose single batch is an `Err`, for proving that the guard does + /// not let an over-budget verdict mask a real error from the child. + #[derive(Debug)] + struct AlwaysErrExec { + properties: Arc, + schema: SchemaRef, + } + + impl AlwaysErrExec { + fn new() -> Self { + let schema = test_schema(); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::clone(&schema)), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { properties, schema } + } + } + + impl DisplayAs for AlwaysErrExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + match t { + DisplayFormatType::Default + | DisplayFormatType::Verbose + | DisplayFormatType::TreeRender => write!(f, "AlwaysErrExec"), + } + } + } + + impl ExecutionPlan for AlwaysErrExec { + fn name(&self) -> &str { + "AlwaysErrExec" + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if !children.is_empty() { + return internal_err!( + "AlwaysErrExec expected no children, got {}", + children.len() + ); + } + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + let schema = self.schema(); + let stream = futures::stream::once(async { + Err(DataFusionError::Execution("boom".to_string())) + }); + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } + + fn partition_statistics( + &self, + _partition: Option, + ) -> Result> { + Ok(Arc::new(Statistics::new_unknown(&self.schema))) + } + } + + #[tokio::test] + async fn passes_batches_through_when_under_budget() { + let _g = oom_guard::test_support::ArmedGuard::acquire(); + oom_guard::arm(0); // unset limit: never trips + + let input = test_input(); + let guard = Arc::new(MemoryGuardExec::new(Arc::clone(&input))); + + assert_eq!(guard.schema(), input.schema(), "schema must pass through"); + + let stream = guard.execute(0, Arc::new(TaskContext::default())).unwrap(); + let batches = collect(stream).await.expect("under budget: must succeed"); + + assert_eq!(batches.len(), 1); + assert_eq!( + batches[0].num_rows(), + 3, + "batches must pass through unchanged" + ); + } + + #[tokio::test] + async fn fails_with_resources_exhausted_when_over_budget() { + // `ArmedGuard` restores the limit and the balance on drop, including on an + // early return from a failed assertion below. + let _g = oom_guard::test_support::ArmedGuard::acquire(); + // The tracking allocator is not installed in this (unit) test binary, so the + // real balance is always zero here. Drive it explicitly to a value far above + // the limit, making the over-budget condition deterministic. + oom_guard::test_support::set_balance_for_test(64 * 1024 * 1024); + oom_guard::arm(1); + + let guard = Arc::new(MemoryGuardExec::new(test_input())); + let stream = guard.execute(0, Arc::new(TaskContext::default())).unwrap(); + let result = collect(stream).await; + + let err = result.expect_err("over budget: the stream must fail"); + assert!( + matches!(err, DataFusionError::ResourcesExhausted(_)), + "must be ResourcesExhausted so the task failure is retriable, got: {err}" + ); + } + + #[tokio::test] + async fn checks_the_budget_before_every_batch_not_just_once() { + // `ArmedGuard` restores the limit and the balance on drop, including on an + // early return from a failed assertion below. + let _g = oom_guard::test_support::ArmedGuard::acquire(); + // Start comfortably under a 1 MiB limit; `StepBudgetExec` raises the balance to + // 64 MiB of its own accord, part-way through the stream. + oom_guard::test_support::set_balance_for_test(0); + oom_guard::arm(1024 * 1024); + + let input: Arc = + Arc::new(StepBudgetExec::new(64 * 1024 * 1024)); + let guard = Arc::new(MemoryGuardExec::new(input)); + let mut stream = guard.execute(0, Arc::new(TaskContext::default())).unwrap(); + + // The first batch was produced while the balance was still under the limit, so + // it must arrive, and arrive with its own exact contents -- not some other + // batch's, and not merely "a" batch. A one-shot guard that checked only once at + // `execute()` time would also pass this assertion, which is why the next one + // matters. + let first = stream + .next() + .await + .expect("stream ended before yielding any batch") + .expect("first batch: budget was under the limit when it was produced"); + let values = first + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + values.values(), + &[1, 2, 3], + "the first batch must arrive with its own exact contents, in order" + ); + + // Producing the batch above already pushed the balance to 64 MiB, over the + // 1 MiB limit. A per-batch check must catch this on the very next item instead + // of letting the remaining two batches through. A one-shot guard -- which + // checked the budget once before the stream ever started, and never again -- + // would instead let all three batches pass, and this is the assertion that + // catches that. + let second = stream + .next() + .await + .expect("stream ended instead of surfacing the over-budget error"); + let err = second.expect_err( + "the balance is now over the limit: the guard must fail, not keep streaming", + ); + assert!( + matches!(err, DataFusionError::ResourcesExhausted(_)), + "must be ResourcesExhausted so the task failure is retriable, got: {err}" + ); + } + + #[tokio::test] + async fn a_child_error_is_not_masked_by_an_over_budget_check() { + let _g = oom_guard::test_support::ArmedGuard::acquire(); + // Over budget for the entire test: this is what previously made + // `check_budget().and(item)` discard the child's own error. + oom_guard::test_support::set_balance_for_test(64 * 1024 * 1024); + oom_guard::arm(1); + + let input: Arc = Arc::new(AlwaysErrExec::new()); + let guard = Arc::new(MemoryGuardExec::new(input)); + let stream = guard.execute(0, Arc::new(TaskContext::default())).unwrap(); + let result = collect(stream).await; + + let err = result.expect_err("the child's error must still surface"); + assert!( + matches!(&err, DataFusionError::Execution(msg) if msg == "boom"), + "an over-budget verdict must not mask the child's own error \ + (a non-retriable failure must not be misreported as retriable), got: {err}" + ); + } + + #[test] + fn preserves_input_plan_properties() { + let input = test_input(); + let guard = MemoryGuardExec::new(Arc::clone(&input)); + + // Inserting the guard must not perturb the plan: same partitioning, same + // ordering, same emission type, same boundedness. If it did, it would change + // query semantics or the scheduler's stage planning. + // + // Identity, not equality, is the assertion: `properties()` must hand back the + // input's `PlanProperties` object itself. That is both the strongest possible + // statement of transparency (every field is the same field, including ones a + // structural comparison would miss) and the only sound one here -- + // `PartialEq for Partitioning` in DataFusion models "satisfies", not structural + // equality, so `UnknownPartitioning` (which this input has) never compares equal + // even to itself, and `assert_eq!` on it would fail for a perfectly transparent + // node. + assert!( + Arc::ptr_eq(guard.properties(), input.properties()), + "the guard must return the input's PlanProperties unchanged" + ); + + // Belt and braces: the fields the plan's correctness actually rests on, compared + // with the operators that do have meaningful equality here. + assert_eq!( + guard.properties().output_partitioning().partition_count(), + input.properties().output_partitioning().partition_count() + ); + assert_eq!( + guard.properties().output_ordering(), + input.properties().output_ordering() + ); + assert_eq!( + guard.properties().emission_type, + input.properties().emission_type + ); + assert_eq!( + guard.properties().boundedness, + input.properties().boundedness + ); + } +} diff --git a/ballista/executor/src/memory_pools/mod.rs b/ballista/executor/src/memory_pools/mod.rs new file mode 100644 index 0000000000..14b79adea7 --- /dev/null +++ b/ballista/executor/src/memory_pools/mod.rs @@ -0,0 +1,49 @@ +// 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. + +//! Allocator-backed OOM protection for the executor. +//! +//! Ballista's memory accounting relies on voluntary `MemoryPool` reservations, which +//! miss allocations made by Arrow buffers, join scratch space, and expression +//! kernels. This module tracks the bytes the global allocator actually hands out and +//! uses that signal in two layers: +//! +//! - a cooperative gate that rejects pool growth (so DataFusion spills and retries) +//! once real usage plus the request would exceed the budget, and +//! - a last-resort circuit breaker that fails a single task with a retriable error +//! rather than letting the process be OOM-killed. +//! +//! The allocator itself only *tracks*: unwinding out of a global allocator is +//! undefined behaviour, so enforcement happens at safe points (pool growth and plan +//! poll boundaries) instead. + +mod guard_exec; +pub mod oom_guard; +mod real_usage_pool; + +pub use guard_exec::MemoryGuardExec; +pub use real_usage_pool::RealUsagePool; + +/// Serializes every test in this crate that reads or writes the process-global limit +/// and balance owned by [`oom_guard`], wherever in the module tree it lives. +/// +/// Those two atomics are shared mutable state for the whole test binary, so a mutex +/// private to one module would not serialize it against a test in another module +/// writing the same globals. See `oom_guard::test_support::ArmedGuard`, which takes +/// this lock and restores both values on drop. +#[cfg(test)] +pub(crate) static GLOBAL_STATE_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); diff --git a/ballista/executor/src/memory_pools/oom_guard.rs b/ballista/executor/src/memory_pools/oom_guard.rs new file mode 100644 index 0000000000..1ec899e2f1 --- /dev/null +++ b/ballista/executor/src/memory_pools/oom_guard.rs @@ -0,0 +1,1154 @@ +// 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. + +//! Real-usage accounting for the executor's global allocator. +//! +//! [`AccountingAllocator`] wraps the real global allocator and tracks the bytes it +//! hands out in a process-global balance. It performs **no enforcement**: unwinding +//! out of a `GlobalAlloc` is undefined behaviour, so the allocator only ever updates +//! a counter. Callers enforce at safe points by calling [`check_budget`] -- the +//! executor does so on every `poll_next` of a running stage, and on every memory-pool +//! growth. +//! +//! # The enforced quantity is *live allocator bytes*, not RSS +//! +//! Enforcement reads a counter of currently-live allocated bytes: every `alloc` adds +//! its layout size, every `dealloc` subtracts it. That is the only quantity that +//! failing a task can actually move. Resident set size (RSS) cannot: an allocator such +//! as mimalloc returns freed pages to the OS lazily (`MADV_FREE`), so RSS stays high +//! for a while after memory is logically free. +//! +//! Gating on RSS would make the guard **latch**: a query trips the limit, DataFusion +//! spills and frees several GB, but RSS has not moved, so every subsequent batch of +//! every task on the executor keeps failing -- and because those failures are +//! retriable, the scheduler re-lands the same tasks on the same executor, where they +//! die again. The guard's own successful spill would take the executor out. So RSS is +//! **advisory only** ([`observe_rss`], which logs and never writes the balance), and +//! the enforced counter decrements the instant memory is freed, which means the guard +//! un-trips within microseconds of a spill and can never latch. +//! +//! # The breaker debounces; the cooperative gate does not +//! +//! [`check_budget`] is the *circuit breaker*, and it only trips after the process has +//! been over budget for several consecutive checks **and** for a minimum wall-clock +//! interval -- see [`Breaker`]. `RealUsagePool::try_grow`, the *cooperative gate*, has +//! no such debounce and rejects immediately. That asymmetry is deliberate and is the +//! whole ordering of the two layers: the gate fires first and cheaply (DataFusion just +//! spills), the breaker is a last resort that fails tasks, so it must give the spill +//! time to land before it starts killing every task on the executor. + +use datafusion::common::{DataFusionError, resources_datafusion_err}; +use log::warn; +use std::alloc::{GlobalAlloc, Layout}; +use std::cell::Cell; +use std::sync::atomic::{AtomicBool, AtomicIsize, AtomicU64, AtomicUsize, Ordering}; +use std::sync::{LazyLock, Mutex}; +use std::time::{Duration, Instant}; + +// The `shard()` thread-local is read from inside `dealloc`, including during thread +// teardown. That is only sound because the key is `const`-initialized and drop-free, so +// std selects its eager, `#[thread_local]`-backed storage, which neither allocates nor +// panics. On a target *without* the `target_thread_local` feature, std falls back to +// lazy OS-key storage that allocates on first access -- from inside `alloc()`, that is +// unbounded re-entrancy into the global allocator. +// +// `cfg(target_thread_local)` is a nightly-only cfg (feature `cfg_target_thread_local`), +// so it cannot be tested on stable; on stable it would simply evaluate to false. Rather +// than invent a fragile probe, this is an allowlist of the targets Ballista ships on, +// all of which have `target_thread_local`. The build fails closed on anything else, +// which is the point: better a compile error than a silent re-entrancy bug. +#[cfg(not(all( + any(target_os = "linux", target_os = "macos", target_os = "windows"), + any(target_arch = "x86_64", target_arch = "aarch64") +)))] +compile_error!( + "the `oom-guard` feature is only supported on linux/macos/windows with x86_64 or \ + aarch64. It requires std's eager `#[thread_local]` TLS storage (the target's \ + `target_thread_local` feature), because the tracking allocator reads a thread-local \ + from inside `dealloc`; on a target without it, std's lazy TLS allocates on first \ + access and re-enters the global allocator. `cfg(target_thread_local)` is unstable, \ + so this allowlist is the enforcement." +); + +/// Number of counter shards. Threads are spread across these so that concurrent +/// allocations rarely contend on the same cache line. +/// +/// Together with the one relaxed `fetch_add` per allocation in [`track`], this is the +/// knob to revisit if allocator overhead shows up in benchmarks. The count trades +/// contention (fewer shards) against the cost of summing them in [`current_balance`] +/// (more shards); the sum is only taken at enforcement points, not on the alloc path. +const SHARD_COUNT: usize = 64; + +/// One shard of the balance, padded to its own cache line so that two threads writing +/// to different shards never false-share. 128 bytes covers the 64-byte lines of x86-64 +/// and the 128-byte prefetch pairing of aarch64. +#[repr(align(128))] +struct Shard(AtomicIsize); + +/// Process-wide outstanding bytes, sharded. The true balance is the sum of all shards. +/// +/// Signed, because a thread may free memory another thread allocated: any individual +/// shard can go arbitrarily negative even though the sum cannot (much) -- and the sum +/// itself may dip below zero transiently while an `alloc`'s `fetch_add` on one shard is +/// in flight against a `dealloc`'s on another. +static SHARDS: [Shard; SHARD_COUNT] = [const { Shard(AtomicIsize::new(0)) }; SHARD_COUNT]; + +/// Hands out shard indices to threads, round-robin, on their first allocation. +static NEXT_SHARD: AtomicUsize = AtomicUsize::new(0); + +/// Enforcement limit in bytes; 0 means unset (never gates). +static LIMIT: AtomicUsize = AtomicUsize::new(0); + +thread_local! { + /// This thread's shard index, assigned on first use. `usize::MAX` means unassigned. + /// + /// A `const`-initialized `Cell` with no drop glue, so `LocalKey::with` cannot panic + /// during thread-local teardown -- which matters, because this is reached from + /// `dealloc` at thread exit. A thread-local that *did* carry a destructor would + /// force std's lazy TLS storage, which allocates -- from inside the global + /// allocator. That guarantee relies on std choosing its eager, + /// `#[thread_local]`-backed storage for this key, which requires the + /// `target_thread_local` feature of the target. That is not a universal guarantee of + /// the language, so it is not left to documentation: the `compile_error!` at the top + /// of this module refuses to build the `oom-guard` feature on any target outside the + /// allowlist of ones known to have it. + /// + /// Because there is no per-thread residue -- every delta lands in a shard + /// immediately -- a thread exiting leaves the balance exactly correct. That is what + /// makes the counter exact, and what removes the need for any periodic correction. + static SHARD_INDEX: Cell = const { Cell::new(usize::MAX) }; +} + +/// Rate-limits a log line to at most one emission per interval. +/// +/// Uses [`Instant`], which is monotonic, rather than [`std::time::SystemTime`]: a +/// backwards wall-clock step (NTP correction, manual clock change) would otherwise +/// suppress the warning for the duration of the step, and a forward step would let an +/// extra line through immediately. `Instant` is not `const`-constructible, so the last +/// emission is held behind a `Mutex` rather than an atomic. That is fine here -- this is +/// reached only from the guard's trip path (`check_budget`, `observe_rss`), never from +/// `GlobalAlloc` -- but it does mean this type must never be used on the alloc/dealloc +/// path itself. +struct RateLimit { + /// The instant of the last emission; `None` means "never yet". + last_emit: Mutex>, +} + +impl RateLimit { + const fn new() -> Self { + Self { + last_emit: Mutex::new(None), + } + } + + /// Whether a log line may be emitted now. Two threads racing here may both emit + /// within one window; that is harmless for an advisory warning. + /// + /// Never panics: a poisoned mutex (only possible if an earlier caller panicked + /// while holding it, which nothing here does) is recovered from rather than + /// propagated. + fn allow(&self, min_interval: Duration) -> bool { + let now = Instant::now(); + let mut last_emit = self.last_emit.lock().unwrap_or_else(|e| e.into_inner()); + if let Some(last) = *last_emit + && now.saturating_duration_since(last) < min_interval + { + return false; + } + *last_emit = Some(now); + true + } +} + +/// How often the guard may log that it is over budget. [`check_budget`] runs once per +/// batch on every task, so an unthrottled log here would flood. +const TRIP_LOG_INTERVAL: Duration = Duration::from_secs(60); + +/// How often the advisory RSS-divergence warning may be logged. +const RSS_LOG_INTERVAL: Duration = Duration::from_secs(60); + +/// Throttles the "over budget" warning emitted when the guard trips. +static TRIP_LOG: RateLimit = RateLimit::new(); + +/// Throttles the advisory "RSS diverges from tracked usage" warning. +static RSS_LOG: RateLimit = RateLimit::new(); + +/// Set once the "this platform does not report RSS" warning has been logged; that fact +/// never changes at runtime, so it is worth saying exactly once. +static RSS_UNAVAILABLE_LOGGED: AtomicBool = AtomicBool::new(false); + +/// How many *consecutive* over-budget observations the breaker requires before it trips. +/// +/// The trade-off. Too low (1, the original behaviour) and a single hungry task that +/// briefly overshoots takes down every other task on the executor with it: the +/// cooperative gate in `RealUsagePool` is at that same instant telling the hungry task +/// to spill, which will drop the balance within microseconds, but tasks B..H poll in +/// the meantime, see the same over-budget balance, and all fail. One greedy task +/// becomes an executor-wide stall. Too high, and a genuinely runaway allocation gets +/// more time to reach the OS OOM killer, which is the outcome the breaker exists to +/// prevent -- and the breaker is the *only* defence against memory the pool never sees. +/// Three is enough to ride out a spill without materially delaying a real trip. +const MIN_CONSECUTIVE_OVER_BUDGET_CHECKS: usize = 3; + +/// How long the process must have been *continuously* over budget before the breaker +/// trips. +/// +/// The count above is not sufficient on its own, because batches arrive at wildly +/// different rates: a stage emitting tiny batches burns three checks in microseconds -- +/// far less time than a spill needs to complete -- so the count alone would reproduce +/// the very fan-out it is meant to prevent. Conversely time alone is not sufficient +/// either: a stage whose batches take seconds would let a single slow observation, taken +/// long ago, trip the breaker on its next poll with no evidence the condition persisted. +/// Requiring *both* means "over budget on every check we took, over a real interval", +/// which is the condition a completed spill would have cleared. +/// +/// 100 ms is comfortably longer than a spill takes to start freeing memory, and short +/// enough that a true runaway is still caught long before the process is killed. +const MIN_TIME_OVER_BUDGET: Duration = Duration::from_millis(100); + +/// The process-global breaker state. +/// +/// In the crate's unit-test build the thresholds are relaxed to "trip on the first +/// over-budget check", so that the tests of *other* properties (`guard_exec`'s per-batch +/// checking, for instance) do not have to sleep. The debounce itself is not tested +/// through this static: [`Breaker`] takes its thresholds and its clock as arguments, and +/// the tests below drive a local instance with explicit values. The production thresholds +/// are exercised end to end in `tests/oom_guard_alloc.rs`, which links this library +/// without `cfg(test)`. +static BREAKER: Breaker = Breaker::new( + if cfg!(test) { + 1 + } else { + MIN_CONSECUTIVE_OVER_BUDGET_CHECKS + }, + if cfg!(test) { + Duration::ZERO + } else { + MIN_TIME_OVER_BUDGET + }, +); + +/// The instant the process started tracking, used as the epoch for [`now_nanos`]. +/// +/// [`Instant`] is monotonic (unlike `SystemTime`, which an NTP correction can step +/// backwards, and which would then either suppress a trip indefinitely or trip it +/// early). It is not `const`-constructible, hence the `LazyLock`; initialising it does +/// not allocate, and nothing on this path is reachable from `GlobalAlloc` anyway. +static PROCESS_START: LazyLock = LazyLock::new(Instant::now); + +/// Monotonic nanoseconds since [`PROCESS_START`], never zero (zero is [`Breaker`]'s +/// "not currently over budget" sentinel). Allocation-free and panic-free. +fn now_nanos() -> u64 { + let elapsed = PROCESS_START.elapsed().as_nanos(); + u64::try_from(elapsed).unwrap_or(u64::MAX).saturating_add(1) +} + +/// Hysteresis for the circuit breaker: it trips only once the process has been over +/// budget for [`Breaker::min_consecutive`] checks in a row **and** for at least +/// [`Breaker::min_duration`]. +/// +/// The single most important property here is the **reset**: the moment any check +/// anywhere in the process observes the balance back under the limit, the streak is +/// discarded. A successful spill therefore un-arms the breaker, which is exactly the +/// hand-off between the two layers -- the cooperative gate spills, the balance drops, +/// and the breaker never fires. +/// +/// The counter is process-global and shared by every task's stream: three checks may +/// come from three different threads. That is intended. The question the breaker asks is +/// "has the *process* stayed over budget", not "has this one task seen it three times". +struct Breaker { + /// Consecutive over-budget observations required to trip. + min_consecutive: usize, + /// Minimum continuous time over budget required to trip. + min_duration: Duration, + /// Length of the current over-budget streak; reset to 0 by any under-budget check. + consecutive: AtomicUsize, + /// [`now_nanos`] of the first observation in the current streak; 0 means "no streak". + over_since_nanos: AtomicU64, +} + +impl Breaker { + const fn new(min_consecutive: usize, min_duration: Duration) -> Self { + Self { + min_consecutive, + min_duration, + consecutive: AtomicUsize::new(0), + over_since_nanos: AtomicU64::new(0), + } + } + + /// Feed one observation in; returns whether the breaker should trip *now*. + /// + /// Allocation-free: a handful of relaxed atomic operations, called once per batch. + fn observe(&self, over_budget: bool, now_nanos: u64) -> bool { + if !over_budget { + // The spill worked (or the load simply passed). Forget the streak entirely. + self.reset(); + return false; + } + + let count = self + .consecutive + .fetch_add(1, Ordering::Relaxed) + .saturating_add(1); + + // Stamp the start of the streak if this observation opened it. Racing threads + // both see a start; whichever lands first wins, and the loser reads it back. + let started = match self.over_since_nanos.compare_exchange( + 0, + now_nanos, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => now_nanos, + Err(existing) => existing, + }; + let elapsed = u128::from(now_nanos.saturating_sub(started)); + + count >= self.min_consecutive && elapsed >= self.min_duration.as_nanos() + } + + /// Discard any in-progress streak. + fn reset(&self) { + self.consecutive.store(0, Ordering::Relaxed); + self.over_since_nanos.store(0, Ordering::Relaxed); + } + + /// The current streak length (test observability). + #[cfg(test)] + fn streak(&self) -> usize { + self.consecutive.load(Ordering::Relaxed) + } +} + +/// Set the enforcement limit in bytes. `0` disables enforcement (tracking continues). +/// +/// `pub` (rather than `pub(crate)`) because the `oom_guard_alloc` integration test binary +/// arms the guard directly. +pub fn arm(limit_bytes: usize) { + LIMIT.store(limit_bytes, Ordering::Relaxed); + // A new limit invalidates any streak accumulated against the old one. + BREAKER.reset(); +} + +/// This thread's shard, assigned round-robin on first use and then stable. +/// +/// Never allocates and never panics: the thread-local is a `const`-initialized, drop-free +/// `Cell`, and the index is masked into range. +#[inline] +fn shard() -> &'static AtomicIsize { + let index = SHARD_INDEX.with(|cell| { + let existing = cell.get(); + if existing != usize::MAX { + return existing; + } + let assigned = NEXT_SHARD.fetch_add(1, Ordering::Relaxed) % SHARD_COUNT; + cell.set(assigned); + assigned + }); + &SHARDS[index % SHARD_COUNT].0 +} + +/// The exact signed balance: the sum of every shard. +fn raw_balance() -> isize { + SHARDS.iter().fold(0isize, |acc, shard| { + acc.wrapping_add(shard.0.load(Ordering::Relaxed)) + }) +} + +/// Current process-wide tracked usage in bytes (never reported negative). +/// +/// This is the sum of the shards: bytes handed out by the global allocator and not yet +/// returned to it. It is *exact* -- every `alloc`/`dealloc` lands in a shard +/// immediately, so no thread carries un-flushed residue and no thread's exit leaves any +/// behind. Reading it is `SHARD_COUNT` relaxed loads, taken only at enforcement points. +/// +/// Note what this is *not*: it is not RSS. It counts requested `Layout` bytes, so it +/// excludes allocator metadata and slack, and it drops the moment memory is freed even +/// though the pages may stay resident. Enforcing on it is deliberate -- see the module +/// docs -- because it is the quantity that failing or spilling a task can actually move. +/// +/// `pub` (rather than `pub(crate)`) because `RealUsagePool::new` installs it as the +/// cooperative gate's balance source and the `oom_guard_alloc` integration test binary +/// reads it directly. +pub fn current_balance() -> usize { + raw_balance().max(0) as usize +} + +/// Fail with a retriable `ResourcesExhausted` if tracked live allocator usage has been +/// over the limit for long enough to trip the breaker. Callers invoke this at safe +/// points -- never from inside the allocator. +/// +/// Two things keep this from taking out the executor when a single task overshoots: +/// +/// - **Hysteresis.** One over-budget observation is not enough; see [`Breaker`]. The +/// cooperative gate (`RealUsagePool::try_grow`) is rejecting growth from the very first +/// over-budget byte, so by the time the breaker's threshold is met, a spill has had its +/// chance and failed to bring usage down. +/// - **No latching.** The balance decrements on every `dealloc`, so a spill that frees +/// memory both resets the streak and puts the next check back under the limit. +/// +/// `pub` (rather than `pub(crate)`) because the `oom_guard_alloc` integration test binary +/// drives it directly. +pub fn check_budget() -> Result<(), DataFusionError> { + let balance = raw_balance(); + let limit = LIMIT.load(Ordering::Relaxed); + + // Allocation-free on the hot path: relaxed atomic loads, one monotonic clock read, + // and nothing allocated unless the breaker actually trips (the error string). + if !BREAKER.observe(over_budget(balance, limit), now_nanos()) { + return Ok(()); + } + + // Logged here, where the trip is *decided*, and rate-limited: `MemoryGuardExec` calls + // this once per batch on every running task, so an unthrottled line would flood the + // log exactly when the executor is under stress. + if TRIP_LOG.allow(TRIP_LOG_INTERVAL) { + warn!( + "Ballista OOM guard is over budget: tracked live allocator usage {} bytes \ + exceeds the limit of {limit} bytes, and has done so continuously for long \ + enough that spilling has not rescued it. Tasks are being failed with a \ + retriable ResourcesExhausted. (Rate-limited to one line per {}s.)", + balance.max(0), + TRIP_LOG_INTERVAL.as_secs() + ); + } + Err(budget_error(balance, limit)) +} + +/// Record an advisory RSS sample; **never** writes the enforced balance. +/// +/// `rss_bytes` is `None` when the platform cannot report process memory. The executor +/// calls this from a periodic ticker purely for observability: it is what lets an +/// operator staring at a storm of `ResourcesExhausted` tell "this query genuinely needs +/// more memory" (RSS tracks the balance) from "the allocator is hoarding freed pages" +/// (RSS stays high while the balance has dropped). +pub(crate) fn observe_rss(rss_bytes: Option) { + let Some(rss) = rss_bytes else { + if !RSS_UNAVAILABLE_LOGGED.swap(true, Ordering::Relaxed) { + warn!( + "Ballista OOM guard: this platform does not report process memory \ + (memory_stats() returned None), so the advisory RSS check is disabled. \ + Enforcement is unaffected -- it reads the allocator's tracked balance." + ); + } + return; + }; + let tracked = current_balance(); + let limit = LIMIT.load(Ordering::Relaxed); + if rss_diverges(rss, tracked, limit) && RSS_LOG.allow(RSS_LOG_INTERVAL) { + warn!( + "Ballista OOM guard: process RSS ({rss} bytes) has climbed above the \ + configured limit ({limit} bytes) even though tracked live allocator usage \ + ({tracked} bytes) has not. The allocator is most likely holding on to freed \ + pages (mimalloc releases them lazily) or otherwise reserving memory the \ + guard cannot see -- this is real memory that enforcement cannot reclaim by \ + failing tasks, because it gates on the tracked value, not RSS. \ + (Rate-limited to one line per {}s.)", + RSS_LOG_INTERVAL.as_secs() + ); + } +} + +/// Only meaningful once a limit is set: the guard enforces on `tracked`, so the +/// dangerous state is RSS above the limit while `tracked` is below it -- real memory +/// the guard cannot see, and cannot reclaim by failing tasks. Comparing RSS to +/// `tracked` alone would warn on every healthy executor, because RSS legitimately +/// includes the binary's text and data, thread stacks, and allocator reserves, none +/// of which the tracking allocator ever sees. +/// +/// Also deliberately silent when *both* `rss` and `tracked` are over `limit`: in that +/// case [`check_budget`] is already tripping on `tracked` and logging its own warning, +/// so a second line here would be redundant noise rather than new information. +fn rss_diverges(rss: usize, tracked: usize, limit: usize) -> bool { + limit != 0 && rss > limit && tracked <= limit +} + +/// The error a tripped breaker raises, as a pure function of `balance` and `limit`. +/// +/// Split out from [`check_budget`] so that the exact variant -- `ResourcesExhausted`, on +/// which the whole retriability path in `ballista/core/src/error.rs` hinges -- can be +/// unit-tested directly with arbitrary inputs, rather than only through the +/// process-global atomics (which would require the tracking allocator to actually be +/// installed). +fn budget_error(balance: isize, limit: usize) -> DataFusionError { + resources_datafusion_err!( + "Ballista OOM guard: the executor's tracked native memory usage is {} bytes, \ + over the limit of {limit} bytes; failing this task", + balance.max(0) + ) +} + +/// Whether `balance` exceeds `limit`. `limit == 0` means unset, and a negative +/// balance never trips. +/// +/// This is the *instantaneous* condition. It is not on its own sufficient to fail a +/// task -- [`Breaker`] debounces it -- but it is exactly what the cooperative gate acts +/// on with no debounce at all. +fn over_budget(balance: isize, limit: usize) -> bool { + limit != 0 && balance > limit.try_into().unwrap_or(isize::MAX) +} + +/// Record a change of `delta` bytes. Never panics, never allocates, never enforces. +/// +/// One relaxed, usually-uncontended `fetch_add` per allocation. `fetch_add` on an +/// `AtomicIsize` wraps on overflow rather than panicking, which is what a `GlobalAlloc` +/// path requires. +#[inline] +fn track(delta: isize) { + shard().fetch_add(delta, Ordering::Relaxed); +} + +/// Wraps an inner global allocator, tracking the layout bytes it hands out. +/// +/// Tracking only: this never enforces a limit and never unwinds. See [`check_budget`]. +pub struct AccountingAllocator { + inner: A, +} + +impl AccountingAllocator { + /// Wrap `inner`, tracking every allocation it serves. + pub const fn new(inner: A) -> Self { + Self { inner } + } +} + +unsafe impl GlobalAlloc for AccountingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { self.inner.alloc(layout) }; + if !ptr.is_null() { + track(layout.size() as isize); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe { self.inner.dealloc(ptr, layout) }; + track(-(layout.size() as isize)); + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { self.inner.alloc_zeroed(layout) }; + if !ptr.is_null() { + track(layout.size() as isize); + } + ptr + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + let new_ptr = unsafe { self.inner.realloc(ptr, layout, new_size) }; + if !new_ptr.is_null() { + // Only account for a realloc that actually happened. The casts and the + // subtraction cannot overflow: a single allocation cannot exceed + // isize::MAX on any real platform. + track(new_size as isize - layout.size() as isize); + } + new_ptr + } +} + +/// Test-only helpers for driving the process-global limit and balance safely. +/// +/// `LIMIT` and the shard array are process-global, so every test in the crate's unit test +/// binary that touches them -- in this module or in any other -- must serialize +/// against the *same* lock and restore what it changed. Hence a single shared guard +/// here rather than a private mutex per module. +#[cfg(test)] +pub(crate) mod test_support { + use super::{BREAKER, LIMIT, Ordering, SHARDS, raw_balance}; + use crate::memory_pools::GLOBAL_STATE_LOCK; + use std::sync::MutexGuard; + + /// Set the tracked balance directly, for tests that need a deterministic value + /// without routing gigabytes through a real allocator. + /// + /// The production path has no such setter: the balance moves *only* via the + /// allocator's deltas, which is precisely what stops the guard from latching. This + /// collapses the whole balance into shard 0, which is equivalent for every reader + /// (they all sum the shards). + pub(crate) fn set_balance_for_test(bytes: isize) { + for shard in SHARDS.iter() { + shard.0.store(0, Ordering::Relaxed); + } + SHARDS[0].0.store(bytes, Ordering::Relaxed); + } + + /// Holds [`GLOBAL_STATE_LOCK`] and restores [`LIMIT`], the balance, and the breaker's + /// streak on drop -- including on an early return from a failed assertion. Without + /// this, a panic partway through a test that armed the guard would leave the global + /// limit set, the balance raised, and a stale over-budget streak in place for every + /// test that runs afterwards in the same binary. + pub(crate) struct ArmedGuard { + #[allow(dead_code)] + lock: MutexGuard<'static, ()>, + balance: isize, + } + + impl ArmedGuard { + /// Acquire the shared lock, recovering from a poisoning left by an earlier panic. + pub(crate) fn acquire() -> Self { + let lock = GLOBAL_STATE_LOCK.lock().unwrap_or_else(|e| e.into_inner()); + BREAKER.reset(); + Self { + lock, + balance: raw_balance(), + } + } + } + + impl Drop for ArmedGuard { + fn drop(&mut self) { + LIMIT.store(0, Ordering::Relaxed); + set_balance_for_test(self.balance); + BREAKER.reset(); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::memory_pools::oom_guard::test_support::{ + ArmedGuard, set_balance_for_test, + }; + + #[test] + fn over_budget_is_true_only_when_a_limit_is_set_and_exceeded() { + assert!(!over_budget(100, 0), "an unset limit never trips"); + assert!(!over_budget(100, 200), "under the limit"); + assert!( + !over_budget(200, 200), + "at the limit: strictly greater is required" + ); + assert!(over_budget(201, 200), "over the limit"); + assert!(!over_budget(-5, 200), "a negative balance never trips"); + } + + #[test] + fn check_budget_is_ok_unless_a_limit_is_exceeded() { + let _g = ArmedGuard::acquire(); + + // An unset limit never gates, whatever the balance. + arm(0); + assert!(check_budget().is_ok(), "an unset limit must never gate"); + + // Nor does a limit far above any plausible balance. + arm(usize::MAX); + assert!(check_budget().is_ok(), "a huge limit must not gate"); + } + + #[test] + fn budget_error_is_a_resources_exhausted() { + // The task-failure retriability logic in `ballista/core/src/error.rs` depends on + // this exact variant, so assert the variant, not merely that some error came back. + let err = budget_error(201, 200); + assert!( + matches!(err, DataFusionError::ResourcesExhausted(_)), + "must be ResourcesExhausted, got: {err:?}" + ); + } + + /// A local breaker with the production thresholds, driven by an explicit clock: no + /// process globals, no sleeping, no dependence on how fast the test host runs. + fn breaker() -> Breaker { + Breaker::new(MIN_CONSECUTIVE_OVER_BUDGET_CHECKS, MIN_TIME_OVER_BUDGET) + } + + /// Nanoseconds, `ms` milliseconds after the (arbitrary) start of a test's timeline. + fn at_ms(ms: u64) -> u64 { + ms * 1_000_000 + 1 // +1: `now_nanos` never returns 0, the "no streak" sentinel. + } + + #[test] + fn the_breaker_never_trips_while_under_budget() { + let breaker = breaker(); + for ms in 0..1_000 { + assert!( + !breaker.observe(false, at_ms(ms)), + "an under-budget process must never trip the breaker" + ); + } + assert_eq!( + breaker.streak(), + 0, + "no over-budget observation was made: the streak must stay at zero" + ); + } + + #[test] + fn the_breaker_does_not_trip_before_the_consecutive_threshold() { + let breaker = breaker(); + // Fewer than MIN_CONSECUTIVE_OVER_BUDGET_CHECKS observations, spread over far + // more than MIN_TIME_OVER_BUDGET: the time condition alone must not be enough. + for i in 0..MIN_CONSECUTIVE_OVER_BUDGET_CHECKS - 1 { + let now = at_ms(1_000 * i as u64); + assert!( + !breaker.observe(true, now), + "trip at observation {i}: fewer than {MIN_CONSECUTIVE_OVER_BUDGET_CHECKS} \ + consecutive checks must not trip the breaker, however long ago the first \ + one was -- this is what stops one hungry task failing every other task \ + on the executor while its own spill is still in flight" + ); + } + } + + #[test] + fn the_breaker_does_not_trip_before_the_time_threshold() { + let breaker = breaker(); + // Plenty of consecutive checks, but all inside MIN_TIME_OVER_BUDGET: a stage + // emitting tiny batches burns the count in microseconds, which is far less time + // than the cooperative gate's spill needs to land. + for i in 0..50 { + assert!( + !breaker.observe(true, at_ms(0) + i * 1_000), + "50 checks within a microsecond must not trip a 100ms debounce" + ); + } + // Once the interval has genuinely elapsed, the same streak trips. + assert!( + breaker.observe(true, at_ms(MIN_TIME_OVER_BUDGET.as_millis() as u64)), + "over budget for the full interval, over many checks: the breaker must trip" + ); + } + + #[test] + fn the_breaker_trips_once_both_thresholds_are_met() { + let breaker = breaker(); + let mut tripped = None; + for i in 0..MIN_CONSECUTIVE_OVER_BUDGET_CHECKS { + // One observation per 100 ms, so the time condition is met exactly when the + // count condition is. + let now = at_ms(MIN_TIME_OVER_BUDGET.as_millis() as u64 * i as u64); + if breaker.observe(true, now) { + tripped = Some(i); + break; + } + } + assert_eq!( + tripped, + Some(MIN_CONSECUTIVE_OVER_BUDGET_CHECKS - 1), + "the breaker must trip on exactly the {MIN_CONSECUTIVE_OVER_BUDGET_CHECKS}th \ + consecutive over-budget check, no earlier and no later" + ); + } + + /// The property the whole fix rests on: **a successful spill un-arms the breaker.** + /// + /// The cooperative gate rejects a grow, the consumer spills, memory is freed, and the + /// very next check sees the balance back under the limit. That single under-budget + /// observation must wipe the streak -- otherwise the breaker would keep counting + /// across the spill and fail tasks for an over-budget condition that no longer holds. + #[test] + fn a_single_under_budget_observation_resets_the_streak() { + let breaker = breaker(); + + // Build a streak that is one observation short of tripping, over a long interval + // so that the *time* condition is already satisfied. + for i in 0..MIN_CONSECUTIVE_OVER_BUDGET_CHECKS - 1 { + assert!(!breaker.observe(true, at_ms(100 * i as u64))); + } + assert_eq!(breaker.streak(), MIN_CONSECUTIVE_OVER_BUDGET_CHECKS - 1); + + // The spill lands: one check sees the process back under budget. + assert!(!breaker.observe(false, at_ms(1_000))); + assert_eq!( + breaker.streak(), + 0, + "an under-budget observation must reset the streak to zero" + ); + + // Usage climbs again. The streak restarts from scratch -- and, critically, so does + // the clock, so the checks that made up the *old* streak cannot carry the new one + // over the line. + for i in 0..MIN_CONSECUTIVE_OVER_BUDGET_CHECKS - 1 { + assert!( + !breaker.observe(true, at_ms(2_000 + 100 * i as u64)), + "the streak restarted: the breaker must not trip early on the strength of \ + observations taken before the spill" + ); + } + // And only now, having been over budget for a full fresh streak, does it trip. + assert!( + breaker.observe( + true, + at_ms(2_000 + 100 * MIN_CONSECUTIVE_OVER_BUDGET_CHECKS as u64) + ), + "a genuinely sustained over-budget condition must still trip the breaker" + ); + } + + #[test] + fn a_tripped_breaker_keeps_tripping_until_it_sees_under_budget() { + let breaker = breaker(); + for i in 0..MIN_CONSECUTIVE_OVER_BUDGET_CHECKS { + breaker.observe(true, at_ms(100 * i as u64)); + } + // Still over budget, so every subsequent check must keep failing its task: the + // breaker is not a one-shot. + for i in 0..10 { + assert!( + breaker.observe(true, at_ms(1_000 + i)), + "while the process remains over budget the breaker must keep tripping" + ); + } + // ...and stops the instant it does not. + assert!(!breaker.observe(false, at_ms(2_000))); + assert!( + !breaker.observe(true, at_ms(2_001)), + "after a reset, a lone over-budget check must not trip" + ); + } + + #[test] + fn the_balance_is_the_sum_of_every_shard() { + let _g = ArmedGuard::acquire(); + set_balance_for_test(0); + + // Deliberately write to shards other than 0: `current_balance` must sum them all, + // not read a single counter. + SHARDS[1].0.store(1_000, Ordering::Relaxed); + SHARDS[SHARD_COUNT - 1].0.store(2_000, Ordering::Relaxed); + assert_eq!(current_balance(), 3_000); + + // A negative shard cancels a positive one (a thread freeing memory another + // thread allocated), and the reported balance is clamped at zero. + SHARDS[1].0.store(-5_000, Ordering::Relaxed); + assert_eq!(raw_balance(), -3_000); + assert_eq!(current_balance(), 0, "a negative balance reports as zero"); + } + + /// The whole point of the fix: the guard must **un-trip by itself** once memory is + /// freed, with no periodic resync, no RSS sample, and no tick of any kind. + /// + /// A guard that enforced on a sticky ground-truth measurement (process RSS, which an + /// allocator like mimalloc keeps resident long after a `free`) would stay tripped + /// here and keep failing every batch of every task on the executor -- a total outage + /// triggered by its own successful spill. This test would fail against that design. + #[test] + fn a_tripped_guard_un_trips_as_soon_as_memory_is_freed() { + let _g = ArmedGuard::acquire(); + set_balance_for_test(0); + + let allocator = AccountingAllocator::new(std::alloc::System); + let layout = Layout::from_size_align(4 * 1024 * 1024, 8).unwrap(); + + // A 1 MiB limit, and a task that allocates 4 MiB: over budget. + arm(1024 * 1024); + assert!( + check_budget().is_ok(), + "nothing allocated yet: under budget" + ); + + let ptr = unsafe { allocator.alloc(layout) }; + assert!(!ptr.is_null()); + let err = check_budget().expect_err("4 MiB against a 1 MiB limit must trip"); + assert!( + matches!(err, DataFusionError::ResourcesExhausted(_)), + "must be ResourcesExhausted so the task failure is retriable, got: {err}" + ); + + // The task spills: the memory is handed back to the allocator. No resync, no + // ticker, no RSS sample -- the very next check must pass. + unsafe { allocator.dealloc(ptr, layout) }; + assert!( + check_budget().is_ok(), + "freeing the memory must un-trip the guard immediately, with no resync: \ + enforcing on a quantity that a spill cannot move would latch the guard and \ + fail every task on this executor" + ); + assert_eq!( + current_balance(), + 0, + "the freed bytes are gone from the balance" + ); + } + + /// The sharded counter is exact across threads: N threads each allocate and free + /// through the accounting allocator, and once they have all exited the balance is + /// back exactly where it started. + /// + /// This is the test the old per-thread-drift design could not pass: it flushed into + /// the shared balance only every 64 KiB and never flushed on thread exit, so every + /// thread that died left up to +-64 KiB baked in permanently. + #[test] + fn the_balance_returns_to_zero_after_many_threads_alloc_and_free() { + let _g = ArmedGuard::acquire(); + set_balance_for_test(0); + + const THREADS: usize = 16; + const ROUNDS: usize = 8; + // Three differently-sized buffers, churned in a rolling fashion: each is freed + // only after the next has been allocated. Every byte is handed back, but the + // *running* total wanders up and down instead of marching to a tidy zero. + // + // That is what makes this a discriminating test. A design that batched deltas in + // a thread-local cell and only flushed past a threshold would leave each thread + // holding an unflushed tail at exit -- here, 8_000 bytes a thread, 128_000 in + // total -- baked into the shared balance forever, and growing without bound as + // `spawn_blocking` threads churn. Exact per-allocation accounting has no tail. + const A: usize = 96_000; + const B: usize = 32_000; + const C: usize = 8_000; + + std::thread::scope(|scope| { + for _ in 0..THREADS { + scope.spawn(|| { + let allocator = AccountingAllocator::new(std::alloc::System); + let a = Layout::from_size_align(A, 8).unwrap(); + let b = Layout::from_size_align(B, 8).unwrap(); + let c = Layout::from_size_align(C, 8).unwrap(); + for _ in 0..ROUNDS { + let pa = unsafe { allocator.alloc(a) }; + let pb = unsafe { allocator.alloc(b) }; + assert!(!pa.is_null() && !pb.is_null()); + unsafe { allocator.dealloc(pa, a) }; + let pc = unsafe { allocator.alloc(c) }; + assert!(!pc.is_null()); + unsafe { allocator.dealloc(pb, b) }; + unsafe { allocator.dealloc(pc, c) }; + } + }); + } + }); + + assert_eq!( + raw_balance(), + 0, + "every allocation was freed, and every thread has exited: an exact counter \ + must be back at zero, with no per-thread residue left behind" + ); + } + + /// Threads that exit while still holding memory must leave the balance holding + /// exactly those bytes -- no more, no less. + #[test] + fn a_thread_that_exits_leaves_no_residue_of_its_own() { + let _g = ArmedGuard::acquire(); + set_balance_for_test(0); + + let allocator = AccountingAllocator::new(std::alloc::System); + let layout = Layout::from_size_align(8 * 1024, 8).unwrap(); + + // Allocate on a thread that then exits; free on this one. + let ptr = std::thread::scope(|scope| { + scope + .spawn(|| unsafe { allocator.alloc(layout) } as usize) + .join() + .unwrap() + }); + assert_ne!(ptr, 0); + assert_eq!( + raw_balance(), + 8 * 1024, + "the exited thread's live allocation must be accounted exactly" + ); + + unsafe { allocator.dealloc(ptr as *mut u8, layout) }; + assert_eq!( + raw_balance(), + 0, + "freeing from a different thread than the one that allocated must balance out" + ); + } + + #[test] + fn rss_diverges_never_warns_with_no_limit_set() { + // A healthy idle executor: no limit armed, RSS dwarfing a small tracked + // balance because of the binary's own text/data, thread stacks, and allocator + // reserves -- none of which `tracked` ever sees. This must never warn. + let small_tracked = 20 * 1024 * 1024; + let idle_rss = 250 * 1024 * 1024; + assert!( + !rss_diverges(idle_rss, small_tracked, 0), + "an unset limit must never warn, however large the RSS/tracked gap" + ); + } + + #[test] + fn rss_diverges_never_warns_when_both_are_under_the_limit() { + let limit = 1024 * 1024 * 1024; + // Same healthy-idle-executor shape as above, but now with a limit armed: RSS + // is nowhere near the limit, so there is nothing to warn about. + assert!(!rss_diverges(250 * 1024 * 1024, 20 * 1024 * 1024, limit)); + // RSS closer to (but still under) the limit, tracked far under it: still fine. + assert!(!rss_diverges(limit - 1, 1024, limit)); + } + + #[test] + fn rss_diverges_warns_when_rss_exceeds_the_limit_but_tracked_does_not() { + let limit = 1024 * 1024 * 1024; + // This is the hoarding signal the check exists to catch: the guard enforces on + // `tracked`, which is under the limit, so it sees nothing wrong -- but RSS has + // climbed past the limit anyway, meaning real memory the guard cannot reclaim + // by failing tasks. + assert!(rss_diverges(limit + 1, limit, limit)); + assert!(rss_diverges(2 * limit, limit / 2, limit)); + } + + #[test] + fn rss_diverges_does_not_warn_when_both_rss_and_tracked_are_over_the_limit() { + let limit = 1024 * 1024 * 1024; + // Here `check_budget` is already tripping on `tracked` and logging its own + // warning; a second line here would be redundant, not new information. + assert!(!rss_diverges(2 * limit, limit + 1, limit)); + } + + #[test] + fn observe_rss_never_moves_the_enforced_balance() { + let _g = ArmedGuard::acquire(); + set_balance_for_test(1024); + arm(4096); + + // A wildly high RSS sample -- the exact situation that used to re-raise the + // balance and latch the guard -- must be advisory only. + observe_rss(Some(64 * 1024 * 1024 * 1024)); + assert_eq!( + current_balance(), + 1024, + "an RSS sample must never write the enforced balance" + ); + assert!( + check_budget().is_ok(), + "an RSS sample must never be able to trip the guard" + ); + + // The unavailable-platform path must also be inert. + observe_rss(None); + assert_eq!(current_balance(), 1024); + } + + #[test] + fn rate_limit_allows_the_first_event_then_throttles() { + let limiter = RateLimit::new(); + assert!( + limiter.allow(Duration::from_secs(3600)), + "the first event is always allowed" + ); + assert!( + !limiter.allow(Duration::from_secs(3600)), + "a second event inside the window is suppressed" + ); + // A zero-length window never suppresses. + let unthrottled = RateLimit::new(); + assert!(unthrottled.allow(Duration::from_secs(0))); + assert!(unthrottled.allow(Duration::from_secs(0))); + } + + /// Drives `alloc` -> `realloc` (grow) -> `realloc` (shrink) -> `dealloc` directly + /// through a standalone `AccountingAllocator` (not the globally-installed one), so + /// the exact signed delta of every step can be asserted without racing the rest of + /// the test binary's allocator traffic through an intervening `Vec`/`Box`. + #[test] + fn accounting_allocator_tracks_alloc_realloc_dealloc_deltas() { + let _g = ArmedGuard::acquire(); + + let allocator = AccountingAllocator::new(std::alloc::System); + let initial_size = 1024 * 1024; // 1 MiB + let grown_size = 3 * 1024 * 1024; // 3 MiB + let shrunk_size = 256 * 1024; // 256 KiB + + let initial_layout = Layout::from_size_align(initial_size, 8).unwrap(); + let before = raw_balance(); + let ptr = unsafe { allocator.alloc(initial_layout) }; + assert!(!ptr.is_null()); + assert_eq!( + raw_balance() - before, + initial_size as isize, + "alloc must track exactly the requested size" + ); + + // Grow. + let after_alloc = raw_balance(); + let ptr = unsafe { allocator.realloc(ptr, initial_layout, grown_size) }; + assert!(!ptr.is_null()); + assert_eq!( + raw_balance() - after_alloc, + (grown_size - initial_size) as isize, + "a growing realloc must track the positive delta" + ); + + // Shrink. + let grown_layout = Layout::from_size_align(grown_size, 8).unwrap(); + let after_grow = raw_balance(); + let ptr = unsafe { allocator.realloc(ptr, grown_layout, shrunk_size) }; + assert!(!ptr.is_null()); + assert_eq!( + raw_balance() - after_grow, + shrunk_size as isize - grown_size as isize, + "a shrinking realloc must track the negative delta" + ); + + // Dealloc, returning to the starting balance. + let shrunk_layout = Layout::from_size_align(shrunk_size, 8).unwrap(); + unsafe { allocator.dealloc(ptr, shrunk_layout) }; + assert_eq!( + raw_balance(), + before, + "dealloc must return the balance to its starting value" + ); + } + + /// A stub inner allocator that always fails, for exercising the null-return paths + /// of `alloc` and `realloc` -- the one place a wrong `if !ptr.is_null()` produces a + /// permanent counter ratchet (tracking bytes for memory that was never handed out). + struct AlwaysNull; + + unsafe impl GlobalAlloc for AlwaysNull { + unsafe fn alloc(&self, _layout: Layout) -> *mut u8 { + std::ptr::null_mut() + } + + unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) { + unreachable!("nothing should ever hold a pointer from this allocator"); + } + + unsafe fn realloc( + &self, + _ptr: *mut u8, + _layout: Layout, + _new_size: usize, + ) -> *mut u8 { + std::ptr::null_mut() + } + } + + #[test] + fn a_failed_alloc_or_realloc_tracks_nothing() { + let _g = ArmedGuard::acquire(); + + let allocator = AccountingAllocator::new(AlwaysNull); + let layout = Layout::from_size_align(4 * 1024 * 1024, 8).unwrap(); + let before = raw_balance(); + + let ptr = unsafe { allocator.alloc(layout) }; + assert!(ptr.is_null()); + assert_eq!( + raw_balance(), + before, + "a failed alloc must not move the balance" + ); + + let ptr = + unsafe { allocator.realloc(std::ptr::null_mut(), layout, layout.size() * 2) }; + assert!(ptr.is_null()); + assert_eq!( + raw_balance(), + before, + "a failed realloc must not move the balance -- this is the ratchet regression test" + ); + } +} diff --git a/ballista/executor/src/memory_pools/real_usage_pool.rs b/ballista/executor/src/memory_pools/real_usage_pool.rs new file mode 100644 index 0000000000..b6ace9846a --- /dev/null +++ b/ballista/executor/src/memory_pools/real_usage_pool.rs @@ -0,0 +1,316 @@ +// 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. + +//! A [`MemoryPool`] decorator that gates growth on *real* allocator usage. + +use crate::memory_pools::oom_guard; +use datafusion::common::{DataFusionError, resources_datafusion_err}; +use datafusion::execution::memory_pool::{ + MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, +}; +use std::sync::Arc; + +/// Source of the current process-wide live allocator usage, in bytes. Production +/// wiring uses [`oom_guard::current_balance`]; tests inject a controllable value. +type BalanceSource = Arc usize + Send + Sync>; + +/// A [`MemoryPool`] decorator that, on top of the inner pool's tracked-reservation +/// accounting, rejects growth when *real* allocator usage (untracked Arrow, join, and +/// kernel bytes included) plus the requested amount would exceed a process-global +/// ceiling. Returning `ResourcesExhausted` lets DataFusion spill and retry. +/// +/// # What the ceiling means +/// +/// The ceiling is the executor's whole `--memory-pool-size`, and it is compared against +/// the live bytes the global allocator has handed out for the **entire process** -- not +/// just query memory. gRPC buffers, tokio's internals, object-store clients and caches, +/// and the shuffle write path all count against it, because they are all real memory the +/// OOM killer would see. So an executor built with the `oom-guard` feature gates somewhat +/// *earlier* than the same `--memory-pool-size` implies in the default build, where the +/// pool only ever sees explicit reservations. No baseline is subtracted: the number +/// budgets the process, not the queries. +/// +/// Rejection is first-come: once the process is over the ceiling, any `try_grow` is +/// rejected, so every running task spills and real usage drops. A small task can be +/// the one told to spill, but spilling is cheap and correct, and a consumer that +/// cannot spill falls through to a retriable task failure. Because the tracked balance +/// falls as soon as the spilled memory is freed, the gate re-opens immediately -- it does +/// not stay closed waiting for the OS to reclaim pages. +pub struct RealUsagePool { + inner: Arc, + /// Process-global real-usage ceiling in bytes; 0 means unset (no gating). + ceiling: usize, + balance_source: BalanceSource, +} + +impl std::fmt::Debug for RealUsagePool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RealUsagePool") + .field("inner", &self.inner) + .field("ceiling", &self.ceiling) + .finish_non_exhaustive() + } +} + +impl RealUsagePool { + /// Wrap `inner` with the real-usage gate, reading the live allocator balance. + /// A `ceiling` of 0 disables gating. + pub fn new(inner: Arc, ceiling: usize) -> Self { + Self { + inner, + ceiling, + balance_source: Arc::new(oom_guard::current_balance), + } + } + + /// Wrap `inner` with an explicit balance source (test seam). + #[cfg(test)] + fn with_balance_source( + inner: Arc, + ceiling: usize, + balance_source: BalanceSource, + ) -> Self { + Self { + inner, + ceiling, + balance_source, + } + } +} + +impl std::fmt::Display for RealUsagePool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{}(ceiling: {}, inner: {})", + self.name(), + self.ceiling, + self.inner + ) + } +} + +impl MemoryPool for RealUsagePool { + fn name(&self) -> &str { + "real_usage" + } + + 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<(), DataFusionError> { + // Check the real-usage ceiling before delegating, so an over-budget request is + // rejected without speculatively reserving the inner pool. + // + // This rejection is **immediate**: no hysteresis, no debounce, no grace period. + // That is the crux of how the two layers order themselves. This is the + // *cooperative* layer, and it is meant to fire first: rejecting a grow is cheap + // and safe -- DataFusion simply spills the consumer and retries -- so the + // earliest possible rejection is the best one. The *circuit breaker* + // (`oom_guard::check_budget`) is the layer that fails whole tasks, and it is the + // one that debounces, precisely so that the spill this rejection triggers has + // time to bring the balance back down before any task is killed. Adding + // hysteresis here would invert that: the guard would start failing tasks before + // the cheap remedy had even been offered. + if self.ceiling != 0 && additional != 0 { + let real = (self.balance_source)(); + if real.saturating_add(additional) > self.ceiling { + return Err(resources_datafusion_err!( + "Ballista real-usage gate: native usage {real} bytes + requested \ + {additional} bytes exceeds the executor memory budget of {} bytes; \ + spilling or failing this consumer", + self.ceiling + )); + } + } + self.inner.try_grow(reservation, additional) + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } + + fn memory_limit(&self) -> MemoryLimit { + self.inner.memory_limit() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::execution::memory_pool::{GreedyMemoryPool, UnboundedMemoryPool}; + use std::sync::atomic::{AtomicUsize, Ordering}; + + fn fixed_source(bytes: usize) -> BalanceSource { + let cell = Arc::new(AtomicUsize::new(bytes)); + Arc::new(move || cell.load(Ordering::Relaxed)) + } + + fn pool_with_balance( + inner: Arc, + ceiling: usize, + balance: usize, + ) -> Arc { + Arc::new(RealUsagePool::with_balance_source( + inner, + ceiling, + fixed_source(balance), + )) + } + + #[test] + fn under_ceiling_delegates_to_inner() { + let inner: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let pool = pool_with_balance(Arc::clone(&inner), 1024 * 1024, 100); + let reservation = + MemoryConsumer::new("test").register(&(pool as Arc)); + + reservation + .try_grow(1024) + .expect("under the ceiling should succeed"); + assert_eq!(inner.reserved(), 1024, "the grow must reach the inner pool"); + } + + #[test] + fn over_ceiling_rejects_without_reserving_inner() { + // Real usage is already at the ceiling, but the inner pool has plenty of room: + // only the real-usage gate can reject this. + let inner: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let pool = pool_with_balance(Arc::clone(&inner), 1000, 1000); + let reservation = + MemoryConsumer::new("test").register(&(pool as Arc)); + + let err = reservation + .try_grow(1) + .expect_err("over the ceiling should be rejected"); + assert!( + matches!(err, DataFusionError::ResourcesExhausted(_)), + "expected ResourcesExhausted so DataFusion spills, got: {err}" + ); + assert_eq!( + inner.reserved(), + 0, + "a rejected grow must not speculatively reserve the inner pool" + ); + } + + #[test] + fn unset_ceiling_never_gates() { + let inner: Arc = Arc::new(UnboundedMemoryPool::default()); + // ceiling == 0 means unset; even a huge real balance must not gate. + let pool = pool_with_balance(Arc::clone(&inner), 0, usize::MAX / 2); + let reservation = + MemoryConsumer::new("test").register(&(pool as Arc)); + + reservation + .try_grow(4096) + .expect("an unset ceiling must never gate"); + assert_eq!(inner.reserved(), 4096); + } + + #[test] + fn zero_sized_grow_is_never_gated() { + let inner: Arc = Arc::new(GreedyMemoryPool::new(1024)); + let pool = pool_with_balance(Arc::clone(&inner), 1000, 5000); + let reservation = + MemoryConsumer::new("test").register(&(pool as Arc)); + + reservation + .try_grow(0) + .expect("a zero-byte grow must never be rejected"); + } + + #[test] + fn grow_exactly_to_the_ceiling_is_allowed_but_one_byte_over_is_not() { + let inner: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let pool = pool_with_balance(Arc::clone(&inner), 1000, 900); + let reservation = MemoryConsumer::new("test") + .register(&(Arc::clone(&pool) as Arc)); + + // real (900) + additional (100) == ceiling (1000): allowed, strictly-greater is required. + reservation + .try_grow(100) + .expect("growing exactly to the ceiling must be allowed"); + + // real (900) + additional (101) == 1001 > ceiling (1000): rejected. + let over = + pool_with_balance(Arc::new(GreedyMemoryPool::new(1024 * 1024)), 1000, 900); + let reservation = + MemoryConsumer::new("test").register(&(over as Arc)); + assert!( + reservation.try_grow(101).is_err(), + "one byte over the ceiling must be rejected" + ); + } + + #[test] + fn a_huge_balance_saturates_rather_than_overflowing() { + // real + additional would overflow usize; saturating_add must clamp so this is + // rejected, not panic. + let inner: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let pool = pool_with_balance(Arc::clone(&inner), 1000, usize::MAX); + let reservation = + MemoryConsumer::new("test").register(&(pool as Arc)); + + assert!( + reservation.try_grow(4096).is_err(), + "a balance near usize::MAX must saturate and reject, not overflow" + ); + assert_eq!( + inner.reserved(), + 0, + "a rejected grow must not reserve the inner pool" + ); + } + + #[test] + fn shrink_and_accessors_delegate_to_inner() { + let inner: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let pool = pool_with_balance(Arc::clone(&inner), 1024 * 1024, 0); + let reservation = MemoryConsumer::new("test") + .register(&(Arc::clone(&pool) as Arc)); + + reservation.try_grow(4096).unwrap(); + assert_eq!(pool.reserved(), 4096, "reserved() must delegate"); + + reservation.shrink(1024); + assert_eq!(inner.reserved(), 3072, "shrink() must delegate"); + assert_eq!(pool.reserved(), 3072); + + assert!( + matches!(pool.memory_limit(), MemoryLimit::Finite(n) if n == 1024 * 1024), + "memory_limit() must delegate to the inner pool" + ); + } +} diff --git a/ballista/executor/tests/oom_guard_alloc.rs b/ballista/executor/tests/oom_guard_alloc.rs new file mode 100644 index 0000000000..422955f5fd --- /dev/null +++ b/ballista/executor/tests/oom_guard_alloc.rs @@ -0,0 +1,284 @@ +// 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. + +//! End-to-end proof that the installed `AccountingAllocator` moves the tracked +//! balance for real heap allocations, that `check_budget` trips on it, and that a +//! [`MemoryGuardExec`] over that balance fails its stream with `ResourcesExhausted`. +//! +//! This lives in its own integration test binary, deliberately separate from the +//! crate's unit test module, for two reasons: +//! +//! - An integration test links `ballista-executor` as a normal dependency, so +//! `cfg(test)` is *not* set for the library. There is therefore no library-side +//! test allocator to conflict with, and this file is free to declare its own +//! `#[global_allocator]`. +//! - Nothing else in this process allocates concurrently, so asserting against the +//! process-global balance is deterministic here -- unlike in the crate's unit test +//! binary, where the tracking allocator is not installed at all (the balance is +//! simply zero), and dozens of unrelated tests run in parallel. The unit tests there +//! must therefore inject a balance; only here is the real causal chain observable. +//! The tests in this file take [`SERIAL`] so they do not perturb each other. +//! +//! The whole file is gated on the `oom-guard` feature: `ballista_executor::memory_pools` +//! itself only exists when that feature is enabled. + +#![cfg(feature = "oom-guard")] + +use ballista_executor::memory_pools::oom_guard::{ + AccountingAllocator, arm, check_budget, current_balance, +}; +use ballista_executor::memory_pools::{MemoryGuardExec, RealUsagePool}; +use datafusion::arrow::array::Int32Array; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::DataFusionError; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::execution::TaskContext; +use datafusion::execution::memory_pool::{FairSpillPool, MemoryConsumer, MemoryPool}; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::common::collect; +use std::sync::{Arc, Mutex, MutexGuard}; +use std::time::{Duration, Instant}; + +/// Poll `check_budget` until the circuit breaker trips, and return the error it raised. +/// +/// The breaker deliberately debounces: it fails a task only once the process has been +/// over budget for several *consecutive* checks and for a minimum interval, so that a +/// spill triggered by the cooperative gate has a chance to rescue the executor before +/// any task is killed. Nothing resets that streak while the memory below is still held, +/// so this converges. +/// +/// This is the only test binary that sees the production thresholds -- the library's own +/// unit tests are built with `cfg(test)`, where the debounce is relaxed to trip on the +/// first check -- so this loop is also what proves the debounced global path can trip at +/// all. It panics rather than hanging if it cannot. +fn wait_for_breaker_to_trip() -> DataFusionError { + let deadline = Instant::now() + Duration::from_secs(10); + loop { + match check_budget() { + Err(err) => return err, + Ok(()) => { + assert!( + Instant::now() < deadline, + "the memory is still held and the process is still over budget, so \ + the breaker must trip: it never did within 10s" + ); + std::thread::sleep(Duration::from_millis(10)); + } + } + } +} + +#[cfg(feature = "oom-guard")] +#[global_allocator] +static GLOBAL: AccountingAllocator = + AccountingAllocator::new(std::alloc::System); + +/// Serializes the tests in this binary. They share the process-global limit and balance, +/// and `cargo test` runs them on parallel threads, so without this one test's `arm` (or +/// its 64 MiB of held allocations) would perturb the other. +static SERIAL: Mutex<()> = Mutex::new(()); + +/// Holds [`SERIAL`] and disarms the limit on drop, including on an early return from a +/// failed assertion -- so a failure here cannot leave the limit armed for the next test. +struct Serial(#[allow(dead_code)] MutexGuard<'static, ()>); + +impl Serial { + fn acquire() -> Self { + Self(SERIAL.lock().unwrap_or_else(|e| e.into_inner())) + } +} + +impl Drop for Serial { + fn drop(&mut self) { + arm(0); + } +} + +#[test] +fn real_allocations_move_the_balance_and_trip_check_budget() { + let _serial = Serial::acquire(); + // 8 MiB of headroom over the current baseline. `Serial` keeps the other test in + // this binary from running concurrently, so there is no allocator traffic to + // introduce noise. + let headroom = 8 * 1024 * 1024; + arm(current_balance() + headroom); + assert!(check_budget().is_ok(), "should start under budget"); + + // Allocate well past the headroom and hold it, so the balance stays raised. + let mut held: Vec> = Vec::new(); + for _ in 0..64 { + held.push(vec![0u8; 1024 * 1024]); + } + // Touch the data so the allocations cannot be optimized away. + assert_eq!( + held.iter().map(|v| v.len()).sum::(), + 64 * 1024 * 1024 + ); + + let err = wait_for_breaker_to_trip(); + assert!(matches!(err, DataFusionError::ResourcesExhausted(_))); + + // Dropping the allocations must bring the balance back under budget -- immediately, + // with no debounce on the way *out*: the breaker's streak is reset by the very first + // under-budget observation, which is what makes a successful spill un-arm it. + drop(held); + assert!( + check_budget().is_ok(), + "freeing the allocations must credit the balance back" + ); + + arm(0); +} + +/// The cooperative gate, built the way production builds it. +/// +/// Every `RealUsagePool` unit test injects a fake balance through the crate-private +/// `with_balance_source` seam, so the *production* constructor -- `RealUsagePool::new`, +/// the only path that installs `oom_guard::current_balance` as the pool's balance source +/// -- is type-checked but never behaviour-checked there. A `new()` that captured `|| 0` +/// would pass every one of those tests while silently disabling the gate in production. +/// +/// This is the test that would catch that: it builds the pool through `new()`, allocates +/// real heap memory until the *real* allocator balance is over the ceiling, and requires +/// the next `try_grow` to be rejected. It can only live here, in the binary where the +/// tracking allocator is genuinely installed. +#[test] +fn the_production_pool_gates_on_the_real_allocator_balance() { + let _serial = Serial::acquire(); + + // A ceiling 8 MiB above wherever the balance happens to sit right now; the inner pool + // is given a gigabyte, so nothing but the real-usage gate can reject a small grow. + let ceiling = current_balance() + 8 * 1024 * 1024; + let inner: Arc = Arc::new(FairSpillPool::new(1024 * 1024 * 1024)); + let pool: Arc = + Arc::new(RealUsagePool::new(Arc::clone(&inner), ceiling)); + let reservation = MemoryConsumer::new("real-usage-test").register(&pool); + + // Under the ceiling: the grow goes through to the inner pool. + reservation + .try_grow(1024) + .expect("under the ceiling: the grow must be allowed"); + assert_eq!(inner.reserved(), 1024, "the grow must reach the inner pool"); + + // Now allocate 64 MiB for real and hold it, so the balance the *installed* allocator + // tracks genuinely rises past the 8 MiB of headroom. + let mut held: Vec> = Vec::new(); + for _ in 0..64 { + held.push(vec![0u8; 1024 * 1024]); + } + assert_eq!( + held.iter().map(|v| v.len()).sum::(), + 64 * 1024 * 1024, + "touch the data so the allocations cannot be optimized away" + ); + assert!( + current_balance() > ceiling, + "the real allocator balance must now be over the ceiling" + ); + + let err = reservation + .try_grow(1024) + .expect_err("over the ceiling: the production pool must reject the grow"); + assert!( + matches!(err, DataFusionError::ResourcesExhausted(_)), + "must be ResourcesExhausted so DataFusion spills and retries, got: {err}" + ); + assert_eq!( + inner.reserved(), + 1024, + "a rejected grow must not reserve the inner pool" + ); + + // Freeing the memory must let the very same pool grow again: the gate reads the live + // balance, it does not latch. + drop(held); + reservation + .try_grow(1024) + .expect("back under the ceiling: the grow must be allowed again"); +} + +/// The single-batch input the guard node is exercised over. Deliberately tiny: the +/// memory that trips the guard is allocated by the test itself, not by the plan. +fn test_input() -> Arc { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None).unwrap() +} + +/// The real causal chain, end to end: a genuine heap allocation raises the balance the +/// installed `AccountingAllocator` tracks, and a `MemoryGuardExec` over that balance +/// fails its stream with `ResourcesExhausted`. +/// +/// The unit tests in `guard_exec.rs` drive the balance with an injected value, because +/// the crate's unit test binary has no tracking allocator installed. This is the only +/// place the chain from `Vec::new` through the allocator to a failed task can actually +/// be proven, so it is proven here. +#[tokio::test] +async fn a_real_allocation_makes_the_guard_node_fail_its_stream() { + let _serial = Serial::acquire(); + + let ctx = Arc::new(TaskContext::default()); + let guard = Arc::new(MemoryGuardExec::new(test_input())); + + // Under budget to begin with: the node is transparent and the batch passes through. + arm(current_balance() + 8 * 1024 * 1024); + let batches = collect(guard.execute(0, Arc::clone(&ctx)).unwrap()) + .await + .expect("under budget: the batch must pass through"); + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 3); + + // Now really allocate past the headroom, and hold it so the balance stays raised. + let mut held: Vec> = Vec::new(); + for _ in 0..64 { + held.push(vec![0u8; 1024 * 1024]); + } + assert_eq!( + held.iter().map(|v| v.len()).sum::(), + 64 * 1024 * 1024, + "touch the data so the allocations cannot be optimized away" + ); + + // The guard node's stream checks the budget once per batch, and this input has a + // single batch -- one check, which the breaker's debounce (deliberately) will not + // trip on. Ride out the debounce first: the memory is still held, so the streak is + // never reset, and once the breaker has tripped it keeps tripping for as long as the + // process stays over budget -- including on the node's own check below. + let err = wait_for_breaker_to_trip(); + assert!(matches!(err, DataFusionError::ResourcesExhausted(_))); + + let err = collect(guard.execute(0, Arc::clone(&ctx)).unwrap()) + .await + .expect_err("over budget: the guard must fail the stream"); + assert!( + matches!(err, DataFusionError::ResourcesExhausted(_)), + "must be ResourcesExhausted so the task failure is retriable, got: {err}" + ); + + // Freeing the memory must let the very same node stream successfully again -- the + // guard reflects the live balance, it does not latch. + drop(held); + let batches = collect(guard.execute(0, ctx).unwrap()) + .await + .expect("back under budget: the batch must pass through again"); + assert_eq!(batches.len(), 1); +} diff --git a/docs/source/user-guide/deployment/cargo-install.md b/docs/source/user-guide/deployment/cargo-install.md index 9069bc0d72..94c83df39d 100644 --- a/docs/source/user-guide/deployment/cargo-install.md +++ b/docs/source/user-guide/deployment/cargo-install.md @@ -72,3 +72,18 @@ When the `spark-compat` feature is enabled, additional functions like `sha1`, `e > **Note:** The `spark-compat` feature provides Spark-compatible expressions and functions only, not full Apache Spark API compatibility. For more details about Spark-compatible functions, see [Spark-Compatible Functions](../spark-compatible-functions.md). + +### OOM Guard (Executor Only) + +To enable the executor's allocator-backed OOM guard, install with the `oom-guard` feature: + +```bash +cargo install --locked --features oom-guard ballista-executor +``` + +This installs a tracking global allocator, so it is only available in +executor binaries built with the feature enabled and has no effect on the +scheduler or CLI. See [Tuning Guide: OOM Guard](../tuning-guide.md#oom-guard-experimental) +for what it does and how to arm it with `--memory-pool-size`. + +> **Note:** This feature is experimental and disabled by default. diff --git a/docs/source/user-guide/tuning-guide.md b/docs/source/user-guide/tuning-guide.md index b7906b166c..3237108567 100644 --- a/docs/source/user-guide/tuning-guide.md +++ b/docs/source/user-guide/tuning-guide.md @@ -91,6 +91,80 @@ The executor refuses to start if the per-task share would round to zero (i.e. When `--memory-pool-size` is not set, the executor behaves as before with no memory pool installed. +## OOM Guard (Experimental) + +The memory pool above only ever sees what DataFusion explicitly reserves +through it. Arrow buffer growth, join build-side scratch space, and several +expression kernels allocate native memory without ever going through a +`MemoryReservation`, so an executor's real memory use can run well above what +the pool believes is checked out. When that happens the process itself can be +OOM-killed by the kernel or by Kubernetes. + +This is more damaging in Ballista than a single failed task. A dead executor +takes every shuffle file it had written with it, so every downstream stage +waiting to read from it fails with a fetch-partition error. That cascades into +stage rollbacks and re-execution, potentially affecting other jobs running +concurrently on the same executor. + +The `oom-guard` feature adds a second, allocator-backed layer of protection on +top of `--memory-pool-size`. It is a **build-time** opt-in: it installs a +global allocator that tracks every byte the executor's process actually hands +out, so it is only available in binaries built with the feature enabled, and +it carries no cost otherwise: + +```sh +cargo build --release --features oom-guard +``` + +A default build has no tracking allocator and no per-allocation overhead. The +feature must be compiled in, and it reuses `--memory-pool-size` as its budget. +If that flag is not set, the guard still tracks usage but never enforces +anything — with `oom-guard` compiled in but no memory pool configured, the +feature is a no-op. + +Once armed, the tracked usage is compared against the same limit at two +points: + +- A cooperative gate on the memory pool itself: once real usage plus a + requested reservation would exceed the budget, the pool rejects the growth + so DataFusion spills instead of continuing to allocate. This is the layer + that does the useful work — most of the time, a spill is all that happens + and no task fails. +- A circuit breaker that checks the same budget between batches of a running + stage. If real usage is already over budget when the cooperative gate could + not prevent it, this layer fails just the one task rather than letting the + process die. A failed task reports a retriable `ResourcesExhausted`, so the + scheduler reschedules it on another attempt, bounded by `--task-max-failures` + (default 4), instead of failing the whole job. + +A few things to keep in mind before enabling this. + +Because the tracked figure is **every live allocation in the process**, not +just query memory, it also includes gRPC buffers, the Tokio runtime, +object-store client caches, and so on. With `oom-guard` enabled, an executor +therefore gates somewhat earlier than the same `--memory-pool-size` value +implies in a default build — queries may start spilling sooner than expected. +This is intentional: the whole point is to budget against what the OOM killer +sees, not just what the query engine reserves. + +Enforcement is also batch-granular, not byte-granular: the circuit breaker +checks the budget between batches, so a single operator that allocates past +the limit within one batch can still bring down the process before the check +runs again. The cooperative gate is the layer that actually prevents this in +most cases; the circuit breaker only shrinks the blast radius once it hasn't. + +Finally, the guard tracks bytes requested from the allocator, not resident set +size (RSS). The two can diverge, because allocators such as mimalloc do not +return freed pages to the OS immediately. The executor logs a warning when RSS +runs substantially above the tracked figure — that is the signal that the +allocator is holding on to freed memory rather than that a spill failed to +help. + +This feature is experimental and disabled by default pending further +benchmark validation. Treat it as a safety net for production clusters that +have seen OOM-killed executors, not as a substitute for sizing +`--memory-pool-size` and `--concurrent-tasks` appropriately. + ## Join Strategy Ballista defaults to **sort-merge join** rather than hash join. This is the