Skip to content

feat: track executor allocations and log memory usage - #2530

Merged
coderfender merged 2 commits into
apache:mainfrom
andygrove:andygrove/ballista-accounting-allocator-pr
Oct 5, 2026
Merged

coderfender merged 2 commits into
apache:mainfrom
andygrove:andygrove/ballista-accounting-allocator-pr

Conversation

@andygrove

@andygrove andygrove commented Oct 3, 2026 •

Copy link
Copy Markdown
Member

Which issue does this PR close?

No linked issue. This brings Comet's executor allocation observability to Ballista, following apache/datafusion-comet#6162 and apache/datafusion-comet#6166, and complements the memory drift observability in apache/datafusion#25651.

Rationale for this change

DataFusion memory pool reservations do not account for every allocation in an executor. Logging outstanding Rust allocator bytes makes that otherwise-untracked memory visible when investigating memory usage or sizing executors.

What changes are included in this PR?

  • Add an AccountingAllocator that wraps the existing allocator without changing its allocation behavior or enforcing a memory limit.
  • Preserve Comet's low-contention accounting: signed process-wide balance, 64 KiB thread-local batching, settlement on thread exit, and a non-panicking fallback during TLS teardown.
  • Install the wrapper around mimalloc in the standalone executor, with a system-allocator fallback. Embedded executors retain their application's global allocator.
  • Log outstanding allocations, current pool reservations, the resolved executor-wide pool budget, and live pool count in MiB at INFO level every 10 seconds using an independent Tokio task. Stop the task when executor startup/execution returns, including error paths.
  • Reuse DataFusion's MemoryPool::reserved() to sample live task pools, without wrapping or changing the pool itself. Track pools weakly, deduplicate shared pools, and release the registry lock before calling pool methods.
  • Expose an opt-in start_executor_process_with_memory_metrics entry point; the existing executor API and embedded-executor behavior are unchanged.
  • Add nine deterministic tests covering allocation/free, zeroed and failed allocations, realloc growth/shrink/failure, batching, thread-exit settlement, TLS teardown, and negative-balance reporting.
  • Add eight pool-observation tests covering budget initialization, reservation totals, shared-pool deduplication, pool lifetime, lock safety, vcore-based sizing, custom pools, and session runtime caching.
  • Document what the measurement includes and excludes.

Are there any user-facing changes?

The standalone executor emits a new INFO log line after pool initialization and every 10 seconds:

Ballista executor memory usage: allocated 5412.3 MiB, reserved 3890.0 MiB, pool size 8192.0 MiB (8 live pools)

The value counts requested Rust allocation sizes, not RSS or memory pool reservations. It excludes allocator fragmentation, retained pages, memory mappings, and direct native-library allocations. Per-thread batching makes the value approximate, with less than 64 KiB of unsettled drift per live thread. There are no changes to memory limits or allocator selection.

reserved is the sum of reservations across distinct live task pools. pool size is the resolved total Ballista budget, including auto-sizing, rather than the current reservations or sum of per-task caps. It is shown as unbounded when Ballista pool limiting is disabled.

Related DataFusion work

DataFusion #25651 merged on September 29, 2026. It adds MemoryDriftTracker and PeakRecordingPool::with_drift_tracker, accepting a caller-supplied allocated-byte counter. Those APIs are on DataFusion main but not in Ballista's locked DataFusion 55.1.0 dependency. Its CountingAllocator also lives in the SLT harness and delegates to System, rather than being a reusable allocator backend wrapper.

This PR therefore keeps the Comet-style mimalloc accounting wrapper and reuses the available MemoryPool::reserved() API, without bumping DataFusion or copying the new drift tracker. The upstream tracker can be used for consumer-attributed drift and peak reporting after a dependency upgrade; a new pool wrapper is not needed for the periodic snapshots here.

Validation

  • cargo test --locked --profile ci -p ballista-executor --lib — all 76 tests pass.
  • cargo test --locked --profile ci -p ballista-executor --no-default-features --lib — all 68 tests pass.
  • cargo build --locked --profile ci -p ballista-executor --bin ballista-executor — passes.
  • cargo clippy --locked --profile ci -p ballista-executor --all-targets -- -D warnings — passes.
  • cargo fmt --all -- --check and Prettier 2.7.1 check — pass.
  • Global allocator smoke test: a held 32 MiB allocation raises the counter, and dropping it returns the counter to the baseline within the documented batching tolerance.
  • Executor binary smoke tests verify bounded and unbounded pool-size logging at 10-second intervals and an auto-sized budget matching the resolved initialization value. A separate immediate connection failure exits without hanging the logging or health tasks.

@andygrove
andygrove marked this pull request as draft October 3, 2026 18:20

@milenkovicm milenkovicm left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @andygrove

@coderfender coderfender left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM ! This should give us more visibility into ballista's mem allocation and potentially improve perf. Thank you Andy

@coderfender
coderfender merged commit 43f2952 into apache:main Oct 5, 2026
25 of 26 checks passed
@andygrove
andygrove deleted the andygrove/ballista-accounting-allocator-pr branch October 5, 2026 19:47
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants