Cooperative Budgets: How DataFusion Stops One Operator From Starving the Whole Runtime

TL;DR: DataFusion treats “one poll_next does a lot of work” as a scheduling bug. Every stream is wrapped in a CooperativeStream that spends one unit of Tokio’s 128-unit task budget per record batch and yields when the budget is exhausted. The EnsureCooperative optimizer rule inserts those wrappers automatically at plan leaves and exchange points. Between the N producer tasks and the M output channels of a shuffle, a global gate keeps at least one channel empty, so the exchange is deadlock-free; the memory bound comes from the pool, not the gate.

The starvation problem

Rust’s Stream trait is pull-based: poll_next is called, and the stream either returns Ready(Some(batch)) or Pending. There is no rule that says poll_next must be cheap. An operator that computes an aggregation over a large dataset can spend milliseconds inside a single poll_next call.

On a Tokio runtime that matters twice over:

1. Other tasks starve. Tokio’s work-stealing scheduler can only run other tasks when the current one yields. A query plan with 16 partitions, all running on a 4-core box, serializes into 4 long-running tasks.
2. Cancellation stops working. DataFusion cancels queries by dropping the output stream, which requires the tasks to reach an await point. A task that never yields is uncancellable.

The module doc in datafusion/physical-plan/src/coop.rs:21-30 states the problem plainly:

//! A single call to `poll_next` on a top-level [`Stream`] may potentially perform a lot of work
//! before it returns a `Poll::Pending`. Think for instance of calculating an aggregation over a
//! large dataset.
//!
//! If a `Stream` runs for a long period of time without yielding back to the Tokio executor,
//! it can starve other tasks waiting on that executor to execute them.
//! Additionally, this prevents the query execution from being cancelled.

Note the second clause: this is not just a throughput problem, it is a correctness-of-cancellation problem.

The mechanism: spend budget, then yield

Tokio has a built-in cooperative scheduling primitive. Each task has a budget of 128 units; tokio::task::coop::poll_proceed(cx) returns Poll::Pending and schedules a wakeup if the budget is exhausted, and made_progress() gives one unit back. DataFusion’s default implementation uses it directly:

// datafusion/physical-plan/src/coop.rs:148-163
#[cfg(any(
    datafusion_coop = "tokio",
    not(any(
        datafusion_coop = "tokio_fallback",
        datafusion_coop = "per_stream"
    ))
))]
{
    let coop = std::task::ready!(tokio::task::coop::poll_proceed(cx));
    let value = self.inner.poll_next_unpin(cx);
    if value.is_ready() {
        coop.made_progress();
    }
    value
}

Read the ordering carefully: poll_proceed is called before the inner stream is polled, and made_progress() is called after it returned Ready. That ordering is the whole game. poll_proceed reserves the right to do one more unit of work; if the budget is gone, it returns Pending without polling the inner stream at all, and Tokio will re-poll later. If you instead check the budget after doing the work, you have already done the expensive thing — you can only yield afterwards, which is exactly the “yield too late” bug.

The three-way cfg is worth noting because it shows the design space. The crate declares the possible values at the workspace level:

# Cargo.toml:323-326
'unexpected_cfgs = { level = "warn", check-cfg = [
    'cfg(datafusion_coop, values("tokio", "tokio_fallback", "per_stream"))',

• tokio (the default): the correct poll_proceed/made_progress dance above.
• tokio_fallback: a compromise for runtimes without the real primitive — check has_budget_remaining() before polling, and consume the budget afterwards, “and just assume that that succeeded” (coop.rs:164-185). The comment is honest that this “may have slightly worse performance”.
• per_stream: a completely local budget, no Tokio involvement:

// datafusion/physical-plan/src/coop.rs:114-121
#[cfg(datafusion_coop = "per_stream")]
// Magic value that matches Tokio's task budget value
const YIELD_FREQUENCY: u8 = 128;

with the poll body:

// datafusion/physical-plan/src/coop.rs:187-202
#[cfg(datafusion_coop = "per_stream")]
{
    if self.budget == 0 {
        self.budget = YIELD_FREQUENCY;
        cx.waker().wake_by_ref();
        return Poll::Pending;
    }

    let value = { self.inner.poll_next_unpin(cx) };

    if value.is_ready() {
        self.budget -= 1;
    } else {
        self.budget = YIELD_FREQUENCY;
    }
    value
}

The per_stream variant is subtly different in kind: it keeps a private budget per stream rather than sharing Tokio’s task budget. That means it cannot be starved by a sibling operator, but it also means it does not participate in Tokio’s global fairness — and it does the expensive work before checking (it checks at the top of the poll, which is fine, but it doesn’t reserve).

Two wrappers, one for statics and one for trait objects

// datafusion/physical-plan/src/coop.rs:452-475
pub fn cooperative<T>(stream: T) -> CooperativeStream<T>
where
    T: RecordBatchStream + Unpin + Send + 'static,
{
    CooperativeStream::new(stream)
}

pub fn make_cooperative(stream: SendableRecordBatchStream) -> SendableRecordBatchStream {
    // TODO is there a more elegant way to overload cooperative
    Box::pin(cooperative(RecordBatchStreamAdapter::new(
        stream.schema(),
        stream,
    )))
}

cooperative is generic, so the inner stream’s poll_next is statically dispatched — no vtable, no dynamic dispatch cost per batch. make_cooperative is the escape hatch for Box<dyn RecordBatchStream>, and it pays for that with an adapter layer. The TODO in the source is a standing invitation to do better.

Making it automatic: EnsureCooperative

Wrapping every stream by hand would be unusable. So there is a physical optimizer rule:

// datafusion/physical-optimizer/src/ensure_coop.rs:60-121
impl PhysicalOptimizerRule for EnsureCooperative {
    fn name(&self) -> &str { "EnsureCooperative" }

    fn optimize(&self, plan: Arc<dyn ExecutionPlan>, _config: &ConfigOptions)
        -> Result<Arc<dyn ExecutionPlan>>
    {
        use std::cell::RefCell;
        let ancestry_stack = RefCell::new(Vec::<(SchedulingType, EvaluationType)>::new());

        plan.transform_down_up(
            |plan| {
                let props = plan.properties();
                ancestry_stack.borrow_mut()
                    .push((props.scheduling_type, props.evaluation_type));
                Ok(Transformed::no(plan))
            },
            |plan| {
                ancestry_stack.borrow_mut().pop();
                let props = plan.properties();
                let is_cooperative = props.scheduling_type == SchedulingType::Cooperative;
                let is_leaf = plan.children().is_empty();
                let is_exchange = props.evaluation_type == EvaluationType::Eager;
                ...

The rule walks the plan maintaining a stack of (SchedulingType, EvaluationType) for the ancestors, and wraps a node in CooperativeExec when it is a leaf (a data source) or an exchange point (an eager operator), provided it is not already cooperative and is not already under a cooperative context.

The “already under a cooperative context” check is the interesting bit:

// datafusion/physical-optimizer/src/ensure_coop.rs:94-108
for (scheduling_type, evaluation_type) in ancestry_stack.borrow().iter().rev() {
    // If nearest ancestor is cooperative, we are under a cooperative context
    if *scheduling_type == SchedulingType::Cooperative {
        is_under_cooperative_context = true;
        break;
    // If nearest ancestor is eager, the cooperative context will be reset
    } else if *evaluation_type == EvaluationType::Eager {
        is_under_cooperative_context = false;
        break;
    }
}

A cooperative context ends at an eager operator. That is semantically correct: an eager operator (a repartition, a sort-preserving merge, a buffer) runs its own tasks and drives its own children independently of downstream demand. Making the parent of an eager operator cooperative does nothing for the eager operator’s children, because the eager operator will poll them from a spawned task regardless. So the rule resets at that boundary and starts again below it.

The metadata this relies on lives in PlanProperties:

// datafusion/physical-plan/src/execution_plan.rs:1364-1378
pub enum SchedulingType {
    /// The stream generated by [`execute`](ExecutionPlan::execute) does not actively participate in
    /// cooperative scheduling. ...
    NonCooperative,
    /// The stream generated by [`execute`](ExecutionPlan::execute) actively participates in
    /// cooperative scheduling by consuming task budget when it was able to produce a
    /// [`RecordBatch`].
    Cooperative,
}

and the complementary EvaluationType (execution_plan.rs:1389-1413), which distinguishes demand-driven (Lazy) operators from eager ones.

Coordinating N tasks and M channels

Cooperative scheduling is the intra-task half of the problem. The inter-task half is: how do N producer tasks feed M output partitions without deadlocking, and without letting the exchange buffer the whole dataset?

RepartitionExec is where this lives, and the answer splits in two: the memory pool handles the buffering (by spilling), and the gate handles the deadlock:

• Each output partition has one channel with an unbounded buffer, shared by all N input tasks (in preserve-order mode it instead has N SPSC channels, one per input partition).
• A global Gate counts empty channels. A sender checks empty_channels == 0 before enqueuing and parks itself if the gate is closed (distributor_channels.rs:218-225).
• The guarantee is liveness, not memory: “at least one channel is empty”. A sender may still pile batches into an already-non-empty channel while a sibling stays empty, so the channel buffer is unbounded. The bound on in-flight data is the output partition’s MemoryReservation, and the gate knows nothing about bytes.

The producer side also has to yield, because pulling from the input and partitioning is a lot of work:

// datafusion/physical-plan/src/repartition/mod.rs:2345-2372
// If the input stream is endless, we may spin forever and
// never yield back to tokio.  See
// https://github.com/apache/datafusion/issues/5278.
//
// However, yielding on every batch causes a bottleneck
// when running with multiple cores. See
// https://github.com/apache/datafusion/issues/6290
//
// Thus, heuristically yield after producing num_partition
// batches
if batches_until_yield == 0 {
    tokio::task::yield_now().await;
    batches_until_yield = partitioner.num_partitions();
} else {
    batches_until_yield -= 1;
}

This is a heuristic, and the comments link the two GitHub issues that forced it into existence. Yielding after every batch costs throughput (issue #6290); never yielding hangs unbounded inputs (issue #5278). Yielding every num_partitions batches is the compromise, and the comment explains why that particular multiple is reasonable: “In round robin this is ideal as each input will get a new batch.”

The buffer that is exactly the right size

spawn_buffered is a small function with a subtle comment:

// datafusion/physical-plan/src/common.rs:186-226
pub fn spawn_buffered(mut input: SendableRecordBatchStream, buffer: usize) -> SendableRecordBatchStream {
    // Use tokio only if running from a multi-thread tokio context
    match tokio::runtime::Handle::try_current() {
        Ok(handle) if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread => {
            let mut builder = RecordBatchReceiverStream::builder(input.schema(), buffer);
            let sender = builder.tx();
            builder.spawn(async move {
                // We call `reserve` (which waits until there's room for at least 1 message in the
                // channel buffer) **before** polling from input to ensure we hold a maximum of
                // `buffer` record batches in memory.
                // Polling from input and then calling send() would block when the channel is full
                // so it would essentially hold `buffer` + 1 record batches:
                // * `buffer`: this many elements would live inside the channel, since this is the
                //   channel's capacity
                // * 1 extra RecordBatch which was produced, but there was no room for it in the
                //   channel, so it's being owned by the send() future, which keeps the batch in
                //   memory while it waits for a slot to free up
                while let Ok(permit) = sender.reserve().await {
                    match input.next().await {
                        Some(item) => permit.send(item),
                        None => break,
                    }
                }
                Ok(())
            });
            builder.build()
        }
        _ => input,
    }
}

reserve() before input.next().await is the trick. The naive order (poll, then send) holds buffer + 1 batches: the channel’s buffer plus the one sitting inside the pending send future. With a tokio::sync::mpsc permit held across the poll, the extra batch never materializes. This is the difference between “bounded” and “bounded plus one”, and the plus one is exactly the kind of thing that turns a 64 MiB budget into a 65 MiB OOM.

Also note the runtime_flavor check: on a current-thread runtime there is no point spawning a task, so the input is returned unwrapped.

Panics, cancellation, and common-runtime

The stream builder is where DataFusion handles the awkward parts of running work on a shared runtime.

// datafusion/physical-plan/src/stream.rs:46-55
/// Creates a stream from a collection of producing tasks, routing panics to the stream.
///
/// Note that this is similar to  [`ReceiverStream` from tokio-stream], with the differences being:
///
/// 1. Methods to bound and "detach"  tasks (`spawn()` and `spawn_blocking()`).
///
/// 2. Propagates panics, whereas the `tokio` version doesn't propagate panics to the receiver.
///
/// 3. Automatically cancels any outstanding tasks when the receiver stream is dropped.

The panics matter. If a worker task panics and the panic is swallowed, the query hangs forever waiting for a batch that will never arrive. So the builder merges the receiver with a “check” stream that drains the join set:

// datafusion/physical-plan/src/stream.rs:132-160
let check = async move {
    while let Some(result) = join_set.join_next().await {
        match result {
            Ok(Ok(())) => {}
            Ok(Err(error)) => return Some(Err(error)),
            Err(e) => {
                if e.is_panic() {
                    // resume on the main thread
                    std::panic::resume_unwind(e.into_panic());
                } else {
                    return Some(exec_err!("Non Panic Task error: {e}"));
                }
            }
        }
    }
    None
};

and merges it with the data stream:

// datafusion/physical-plan/src/stream.rs:163-172
let rx_stream = futures::stream::unfold(rx, |mut rx| async move {
    let next_item = rx.recv().await;
    next_item.map(|next_item| (next_item, rx))
});

// Merge the streams together so whichever is ready first
// produces the batch
futures::stream::select(rx_stream, check_stream).boxed()

futures::stream::select polls both and returns whichever is ready first, so an error from a worker surfaces immediately rather than after the remaining workers finish.

datafusion/common-runtime/ wraps Tokio’s task types with DataFusion’s semantics. SpawnedTask aborts on drop:

// datafusion/common-runtime/src/common.rs:108-112
impl<R> Drop for SpawnedTask<R> {
    fn drop(&mut self) {
        self.inner.abort();
    }
}

which is what makes query cancellation work: dropping the output stream drops the SpawnedTasks, which abort the worker tasks, which drop their reservations, which return memory to the pool. The chain of RAII is what makes “cancel a query” a single drop.

spawn_blocking — and where not to use it

DataFusion does use spawn_blocking, but only for genuinely blocking work: JSON parsing (datasource-json/src/source.rs:510), catalog listing (catalog/src/stream.rs:470).

It explicitly does not use it for CPU-bound work:

// datafusion/core/src/lib.rs:644-650
//! Note that DataFusion does not use [`tokio::task::spawn_blocking`] for
//! CPU-bounded work, because `spawn_blocking` is designed for blocking **IO**,
//! not designed CPU bound tasks. Among other challenges, spawned blocking
//! tasks can't yield waiting for input (can't call `await`) so they
//! can't be used to limit the number of concurrent CPU bound tasks or
//! keep the processing pipeline to the same core.

This is a well-argued point and it is the reason the cooperative machinery exists at all. A spawn_blocking task is a thread-pool task: it cannot await, it cannot be cancelled, it cannot participate in Tokio’s budget, and it does not stay on the same core. If you ran a hash aggregation there, you would lose cancellation, lose fairness, and lose the ability to apply backpressure. So CPU-bound operators stay on the async runtime and cooperate.

The memory-table case

MemoryStream is the simplest possible stream: it holds a Vec<RecordBatch> and returns them one at a time, Ready every time:

// datafusion/physical-plan/src/memory.rs:94-100
fn poll_next(
    mut self: std::pin::Pin<&mut Self>,
    _: &mut Context<'_>,
) -> Poll<Option<Self::Item>> {
    if self.index >= self.data.len() {
        return Poll::Ready(None);
    }
    self.index += 1;

It never returns Pending, so on its own it would happily hand out a million batches without ever yielding. The fix is at the source:

// datafusion/datasource/src/memory.rs:79-93
pub fn open(&self, partition: usize, _context: Arc<TaskContext>) -> Result<SendableRecordBatchStream> {
    Ok(Box::pin(cooperative(
        MemoryStream::try_new(
            self.partitions[partition].clone(),
            Arc::clone(&self.projected_schema),
            self.projection.clone(),
        )?
        .with_fetch(self.fetch),
    )))
}

cooperative(...) wraps it, so every 128 batches the task yields to the runtime. LazyMemoryExec goes further and declares itself cooperative at construction time (memory.rs:228-234):

.with_scheduling_type(SchedulingType::Cooperative)

which tells the EnsureCooperative rule that it does not need to wrap it.

The reservation story for memory tables is handled by ReservationStream (stream.rs:745-806), which shrinks the reservation as batches are consumed and frees it entirely at end-of-stream:

// datafusion/physical-plan/src/stream.rs:781-795
Some(Ok(batch)) => {
    self.reservation.shrink(get_record_batch_memory_size(&batch));
    Poll::Ready(Some(Ok(batch)))
}
None => {
    // Stream is done so free the reservation completely
    self.reservation.free();
    // Release the input pipeline's resources.
    let inner_schema = self.inner.schema();
    self.inner = Box::pin(EmptyRecordBatchStream::new(inner_schema));
    Poll::Ready(None)
}

Note the “release the input pipeline’s resources” step: replacing the inner stream with an empty one tears down the upstream chain at end of stream rather than waiting for the outer stream to be dropped. That is what lets a LIMIT 1 over a huge scan release the scan’s memory immediately instead of holding it until the whole plan is dropped.

Trade-offs and failure modes

• Cooperative scheduling costs a poll per batch. Every poll_next on a cooperative stream does an atomic load and a branch. For small batches in a tight loop this is measurable. The per_stream variant exists partly to make that cheaper.
• The budget is a magic number shared with Tokio. YIELD_FREQUENCY = 128 is hardcoded to match Tokio’s internal task budget (coop.rs:118-121). If Tokio changes its budget, DataFusion’s per_stream mode silently drifts out of sync.
• A cooperative stream inside a non-cooperative one is useless. If an operator’s poll_next never returns to the executor, the budget it spends is irrelevant. EnsureCooperative is the safety net, and it depends on operators correctly reporting SchedulingType and EvaluationType — get that wrong and the rule silently skips a subtree.
• Gate-based backpressure is a liveness tool, not a memory tool. It keeps at least one channel empty so senders cannot deadlock; it does not cap a channel’s buffer. Memory is bounded by the pool plus the per-operator reservations — and the channel buffer sits outside both.
• The reserve-before-poll trick is easy to get wrong. A single send() instead of permit.send() puts you back to buffer + 1 batches, and the regression looks like a mysterious 1-batch memory increase.
• spawn_blocking tasks cannot be cancelled once started. SpawnedTask’s doc says it: “Note that if the task was spawned with spawn_blocking, it will only be aborted if it hasn’t started yet” (common-runtime/src/common.rs:28-31). So a blocking task that is stuck in a syscall will outlive the query.
• Cooperative scheduling does not bound memory by itself. It bounds CPU monopolization. Memory is bounded by the memory pool plus the exchange gate plus the per-operator reservations. All three have to line up.

The design is worth stealing: treat “how long can a single poll_next run” as a first-class property of the execution model, encode it in plan metadata, and enforce it with an optimizer rule rather than with discipline. Once the runtime has a notion of a budget, backpressure stops being a memory accounting problem and becomes a scheduling problem — and the scheduler already knows how to stop a task from running forever.