Skip to content

fair_unified memory pool caps the pool total at pool_size / num_consumers instead of each consumer's share #5961

Description

@andygrove

Describe the bug

fair_unified, the default off-heap memory pool, no longer enforces a per-consumer share. Since the DataFusion 53 upgrade (commit 90633dc, #3629) CometFairMemoryPool::try_grow checks the pool-wide total against pool_size / num_consumers:

let used = state.used;
if limit < used + additional {

Before that commit it checked the calling reservation's own size:

let size = reservation.size();
if limit < size + additional {

So the whole pool is now capped at pool_size / num, and registering another consumer anywhere in the task tightens the ceiling on every consumer already running. With three consumers registered, the task as a whole can use a third of its budget. The tuning guide still documents the old behavior: "prevents operators from using more than an even fraction of the available memory (i.e. pool_size / num_reservations)".

Root cause

The commit changed both shrink and try_grow under the same comment, saying DataFusion 53 reorders the reservation's atomic size updates around the pool calls. That is true for shrink and not for try_grow:

  • MemoryReservation::shrink in 52.1.0 called pool.shrink and then updated size; in 53.0.0 it decrements size first (52.1.0, 53.0.0). Comet's shrink check did need to stop reading reservation.size().
  • MemoryReservation::try_grow calls pool.try_grow before adding to size in both versions (52.1.0, 53.0.0). reservation.size() inside the pool was the pre-grow value before the upgrade and still is, so the original per-consumer check was correct and did not need to change.

The DataFusion 54 and 55 upgrades did not touch this code. The regression is in Comet.

Steps to reproduce

Register two consumers on a CometFairMemoryPool with pool_size = 100, grow the first by 40, then try to grow the second by 20. The second consumer has used 0 of its 50-byte share, but the request is denied because state.used + additional = 60 > 50.

Expected behavior

Each consumer is limited to pool_size / num_consumers, as documented, and the pool total is limited by Spark's ledger through the existing acquire call. Restoring reservation.size() + additional in try_grow (keeping the state.used change in shrink) does that, and a unit test over two consumers should pin the per-consumer semantics.

Additional context

CometFairMemoryPool is in native/core/src/execution/memory_pools/fair_pool.rs. The documentation is in docs/source/user-guide/latest/tuning.md. Related to #4576 and the memory management contributor guide in #5933.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

area:memoryMemory pools, reservations, OOM handlingbugSomething isn't workingrequires-triage

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions