[andygrove]

Which issue does this PR close?

• Part of Find a solution to prevent processes being OOM killed due to inaccurate memory accounting #25650.

Rationale for this change

MemoryPool accounting only covers memory that operators explicitly reserve, so a process can be OOM killed while the pool reports plenty of headroom. As a first step, this makes the gap visible: it logs the difference between what pools have reserved and what's actually allocated, and it's on by default in SLT. It only logs and never fails a test, unlike #22626 which was reverted in #22860.

What changes are included in this PR?

• datafusion-execution: MemoryDriftTracker and DriftLoggingPool. The pool wraps any MemoryPool and reports reservation changes to a tracker, which can be shared by many pools. The tracker compares the reserved total against allocated bytes from a caller-supplied Fn() -> usize, since DataFusion doesn't choose the global allocator. It logs at info each time positive drift rises by 64 MB, naming the pool and consumer, and records the peak.
• sqllogictest: a CountingAllocator global allocator, with every test file's pool wrapped against one process-wide tracker. The peak drift is printed at the end of the run. Disable it with --memory-drift false.

Test files run concurrently, so the comparison is process-wide rather than per file. Files that SET datafusion.runtime.memory_limit replace their pool and drop out of the reserved total (6 of 504 files).

Example output from a full local run:

Peak memory drift: drift=84.3 MB allocated=85.4 MB reserved=1176.3 KB pool=order.slt consumer=ExternalSorterMerge[0]

The allocator batches its counts per thread, so SLT runtime is unchanged locally (7s both with and without it).

Are these changes tested?

Yes. There are unit tests for the tracker and pool, plus a doc test. I also ran the full SLT suite locally with drift logging on and off.

Are there any user-facing changes?

New public types in datafusion_execution::memory_pool. No changes to existing APIs.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 94.31438% with 17 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.58%. Comparing base (bdad988) to head (24a789f).

Files with missing lines	Patch %	Lines
datafusion/execution/src/memory_pool/drift.rs	95.00%	7 Missing and 2 partials ⚠️
datafusion/sqllogictest/src/memory_drift.rs	95.00%	1 Missing and 3 partials ⚠️
datafusion/sqllogictest/bin/sqllogictests.rs	76.92%	1 Missing and 2 partials ⚠️
...fusion/execution/src/memory_pool/peak_recording.rs	94.73%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff            @@
##             main   #25651    +/-   ##
========================================
  Coverage   82.57%   82.58%            
========================================
  Files        1142     1144     +2     
  Lines      440861   441156   +295     
  Branches   440861   441156   +295     
========================================
+ Hits       364047   364319   +272     
- Misses      54818    54828    +10     
- Partials    21996    22009    +13     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[adriangb]

Thanks @andygrove. Worked with an agent to review this, overall loooks good! I think making this log-only is a good response to the objection in #22723. The new types in datafusion-execution also look useful for downstream projects such as Comet.

My main concern is the attribution when test files run at the same time. I ran the PR locally:

Run	Peak drift	Allocated	Reserved	Attributed to
Full SLT, default concurrency	60.9 MB	61.5 MB	576.9 KB	push_down_filter_regression.slt / RepartitionExec[1]
Your run (PR description)	84.3 MB	85.4 MB	1176.3 KB	order.slt / ExternalSorterMerge[0]
order.slt only	2.4 MB	2.7 MB	228.4 KB	order.slt / TopK[0] (next run: TopK[1])
Full SLT, --test-threads 1, run 1	43.3 MB	43.5 MB	144.5 KB	push_down_filter_regression.slt / ParquetSink(SerializedFileWriter)
Full SLT, --test-threads 1, run 2	43.6 MB	43.7 MB	144.5 KB	push_down_filter_regression.slt / ParquetSink(SerializedFileWriter)
push_down_filter_regression.slt only	42.9 MB	42.9 MB	0 B	push_down_filter_regression.slt / ParquetSink(SerializedFileWriter)

With the default concurrency, the attributed file is not reliable. The 84 MB peak that your run attributes to order.slt comes from other files, because order.slt alone peaks at 2.4 MB.

With --test-threads 1, the result is stable and correct: two full runs give the same file and approximately the same value as when that file runs alone. A full run takes approximately 65 s on my machine, so this is a mode for investigation, not for each CI run. The output also gives a lead. With a 4 MB log threshold (a local change), the 10M-row COPY in push_down_filter_regression.slt shows the same cycle for each row group:

• While the columns are encoded, allocated is 36–44 MB and ParquetSink(ArrowColumnWriter) reserves 16.7 MB.
• Between row groups, allocated stays at 28–43 MB, but only 144.5 KB is reserved.

Some of this is fixed-size by design (for example, the 10 MB object store write buffer). I did not find the full cause.

My suggestions:

1. Document --test-threads 1 as the way to use this (README suggestion inline).
2. When files run concurrently, say in the summary line that the attribution can be wrong (suggestion inline).
3. Make the log threshold configurable in SLT. At 64 MB, a full run logs nothing.

Also, at SLT scale, most of the drift is memory that the MemoryPool docs exclude on purpose (in-flight batches, Parquet read buffers). A peak is a lead to investigate, not necessarily a bug. The README can say this too.

[adriangb]

This function is called on every grow, try_grow and shrink of every pool that reports to the tracker. For the SLT atomic, that is cheap. But the module docs suggest "allocator statistics", and to get a current value from jemalloc's stats.allocated you must advance epoch through mallctl on each call. That can slow down queries. Can we document this requirement?

Suggested change

	/// Returns the number of bytes currently allocated by the process.
	pub type AllocatedBytesFn = Arc<dyn Fn() -> usize + Send + Sync>;
	/// Returns the number of bytes currently allocated by the process.
	///
	/// This is called on every reservation change (`grow`, `try_grow` and
	/// `shrink`) of every pool that reports to the tracker, so it must be cheap,
	/// e.g. a single atomic load. Do not read allocator statistics that need a
	/// refresh on each call (such as jemalloc's `epoch`).
	pub type AllocatedBytesFn = Arc<dyn Fn() -> usize + Send + Sync>;

[andygrove]

Done in 956f9a0, using your wording.

[adriangb]

Drift is only sampled when a reservation changes. A query that allocates a lot but reserves little is the case this tool must find. That query is only sampled when a different reservation changes, possibly in a different file or after the memory is freed.

A possible solution: make it possible to sample without a reservation change, so that the caller can also sample from a timer or from the allocator flush path. For example:

/// Compare the current allocated bytes with the reserved total now,
/// without a reservation change.
pub fn sample(&self, source: &str) {
    self.observe(source, "", self.reserved());
}

[andygrove]

Added MemoryDriftTracker::sample(source) with a test in 956f9a0. I did not add a timer in the SLT runner yet: timer samples would all be recorded under one source name, so the peak's file and consumer would say even less. Happy to add it in a follow-up if it's useful.

[adriangb]

No test covers the logging logic: a log line at the threshold, and a re-arm after drift falls. The last_logged logic uses a CAS and has an edge case (falling drift lowers the baseline). A test would help here. One option: observe returns bool for "logged", or the tracker keeps a logged_count, and the tests assert on it with with_log_threshold(1000).

[andygrove]

observe now returns true when it logs. New tests in 956f9a0 cover logging at the threshold, re-arming after drift falls (the baseline-lowering case), and negative drift counting as zero.

[adriangb]

The unflushed count of a thread is lost when the thread exits. Tokio blocking-pool threads exit after they are idle, and these threads are used for local file I/O. So the error is not limited to threads * FLUSH_BYTES, and it can grow over a long run. Also, allocated_bytes() clamps at zero, so a negative bias is not visible.

I understand why UNFLUSHED has no destructor (the destructor registration can allocate from inside the allocator). If you keep this design, please correct the comment:

Suggested change

	/// Per-thread count is flushed to [`ALLOCATED`] once it moves this far, so
	/// threads do not contend on one atomic for every allocation. The global count
	/// is therefore accurate to within `threads * FLUSH_BYTES`.
	const FLUSH_BYTES: isize = 256 * 1024;
	/// Per-thread count is flushed to [`ALLOCATED`] once it moves this far, so
	/// threads do not contend on one atomic for every allocation. The global count
	/// is therefore off by up to `FLUSH_BYTES` per live thread. The unflushed
	/// count of a thread that exits (e.g. an idle Tokio blocking thread) is lost,
	/// so this error can grow during a long run.
	const FLUSH_BYTES: isize = 256 * 1024;

[andygrove]

Kept the design and used your comment in 956f9a0.

[comphead]

A cheap way to avoid the loss without a TLS destructor: tokio runs the on_thread_stop hook on every worker and blocking thread just before it exits, outside the allocator, so the runner can flush there.

/// Flush this thread's unflushed count into `ALLOCATED`.
pub fn flush_thread_allocations() {
    let _ = UNFLUSHED.try_with(|unflushed| {
        ALLOCATED.fetch_add(unflushed.replace(0), Ordering::Relaxed);
    });
}

tokio::runtime::Builder::new_multi_thread()
    .enable_all()
    .on_thread_stop(flush_thread_allocations)
    .build()?

In a standalone check with the same counting scheme, 40 spawn_blocking tasks each returned a 100 KB Vec, and their threads then went idle and exited. Without the hook, 21 KB of the 4000 KB was counted. With it, 3919 KB was counted.

Dropping COUNTING and counting from process start would also remove the other source of negative bias, which is memory allocated before enable_memory_drift_logging and freed after it. The PR description reports no runtime change from counting, so --memory-drift false would only need to skip the pool wrapping.

[andygrove]

Done in f79c1f9: flush_thread_allocations is registered as the runtime's on_thread_stop hook, and I removed COUNTING, so counting starts with the process. --memory-drift false now only skips wrapping the pools.

[adriangb]

These 6 files are mostly the spill tests, where the accounting is most important. After the SET, their reservations are not counted, but their allocations are. This makes the drift larger, and a peak can then be attributed to a different file. Is it possible to wrap the new pool too? If not, the README can say that the total is not correct while these files run.

[andygrove]

Fixed in 956f9a0: after each statement, the SLT runner wraps the pool again if a statement replaced it (checked with pool.is::<DriftLoggingPool>()). When aggregate_memory_spill.slt runs alone, queries after its SET datafusion.runtime.memory_limit = '1M' still log with non-zero reserved bytes.

[adriangb]

These fields show which reservation change took the sample. They do not show where the untracked memory came from. When files run concurrently, even the pool can be wrong: a full SLT run attributes the peak to order.slt / ExternalSorterMerge[0] (84 MB), but order.slt alone peaks at 2.4 MB. With --test-threads 1, the pool is correct, but the consumer is still only where the sample was taken. For example, ParquetSink(SerializedFileWriter) is named at the peak, but the total reserved at that moment is 144.5 KB of 43.5 MB allocated. Can the docs say this, so that nobody reads it as blame?

Suggested change

	/// Label of the [`DriftLoggingPool`] that made the observation.
	pub pool: String,
	/// Consumer whose reservation change triggered the observation.
	pub consumer: String,
	/// Label of the [`DriftLoggingPool`] whose reservation change took this
	/// sample. When several pools share a tracker, the untracked memory can
	/// come from any of them.
	pub pool: String,
	/// Consumer whose reservation change took this sample. This shows when
	/// drift was sampled, not what caused it.
	pub consumer: String,

[andygrove]

Done in 956f9a0, using your wording, plus a mention of the new sample source.

[adriangb]

With the default concurrency, this line names a file and consumer that can be wrong (see the table in the review summary). A short hint here tells the reader how to get a correct result:

Suggested change

	if let Some(peak) = memory_drift_tracker().and_then(|t| t.peak_drift()) {
	eprintln!("Peak memory drift: {peak}");
	}
	if let Some(peak) = memory_drift_tracker().and_then(|t| t.peak_drift()) {
	eprintln!("Peak memory drift: {peak}");
	if options.test_threads > 1 {
	eprintln!(
	"Test files ran concurrently, so the file and consumer above can be wrong. \
	Run with --test-threads 1 to attribute drift to one file."
	);
	}
	}

[andygrove]

Added in 956f9a0.

[adriangb]

With --test-threads 1, the results are stable and attributed to the correct file. Two full runs gave push_down_filter_regression.slt at 43.3 MB and 43.6 MB, and that file alone gives 42.9 MB. I suggest that we document this as the way to investigate:

Suggested change

	Test files run concurrently, so the comparison is process-wide: allocated bytes
	across the whole process against reservations summed across all files. Files
	that `SET datafusion.runtime.memory_limit` replace their pool and drop out of
	the reserved total.
	Test files run concurrently by default, so the comparison is process-wide: allocated bytes across the whole process against reservations summed across all files. In this mode, the file in the output can be wrong. To attribute drift to one file, run one file at a time:
	```shell
	cargo test --test sqllogictests -- --test-threads 1
	```
	In both modes, the consumer in the output is the reservation change that took the sample, not necessarily the code that allocated the untracked memory. Most drift at this scale is memory that `MemoryPool` does not track by design (for example, in-flight batches), so a large drift is a lead to investigate, not necessarily a bug.
	Files that `SET datafusion.runtime.memory_limit` replace their pool and drop out of the reserved total.

[andygrove]

Added in 956f9a0. I also removed the sentence about SET files dropping out of the reserved total, because the runner now wraps the replacement pool too.

[adriangb]

With the default 64 MB threshold, this command printed no lines in a full run, with or without --test-threads 1, because the largest drift was 43 MB. With a 4 MB threshold (a local change), the output was useful: it showed each row-group cycle of the COPY in push_down_filter_regression.slt. Can the threshold be an SLT option that is passed to MemoryDriftTracker::with_log_threshold? For example, --memory-drift-log-threshold / SLT_MEMORY_DRIFT_LOG_THRESHOLD.

[andygrove]

Added --memory-drift-log-threshold / SLT_MEMORY_DRIFT_LOG_THRESHOLD (in bytes) in 956f9a0. It is passed to with_log_threshold, and the README has a 4 MB example with --test-threads 1.

[adriangb]

thank you

[comphead]

Thanks @andygrove. A few suggestions inline, plus some small ones:

• allocated_bytes is re-exported from datafusion_sqllogictest but is only used inside memory_drift.rs. CountingAllocator's A parameter is only ever System. Both could go.
• negative_drift_counts_as_zero_for_logging could be two asserts at the start of logs_each_rise_by_the_threshold_and_rearms_after_a_fall. If the tracker() helper took a threshold, it would also replace the hand-built tracker in one test and the Arc::into_inner(..).unwrap() in the other.
• The substrait round-trip engine's run does not call rewrap_replaced_pool, so the SET datafusion.runtime.memory_limit files still drop out of the total with --substrait-round-trip.
• The --memory-drift help says RUST_LOG=info shows each change. Only rises of --memory-drift-log-threshold bytes are logged, and the README uses the narrower RUST_LOG=datafusion_execution::memory_pool=info.

[comphead]

This is close to a copy of PeakRecordingPool in peak_recording.rs: the same pass-through of every MemoryPool method, the same Display passthrough, the same "inner is expected to be empty" caveat, and near-identical delegation and failed-growth tests. Each pass-through wrapper also has to remember to forward any method added to MemoryPool later, because the defaulted ones (register, unregister, memory_limit) silently fall back to the default.

Could PeakRecordingPool report to the tracker instead, e.g. PeakRecordingPool::new(pool).with_drift_tracker(tracker, label)? That removes this type, its Debug/Display impls and the two duplicated tests. DriftLoggingPool::tracker() also has no caller.

[andygrove]

Done in f79c1f9: I removed DriftLoggingPool and added PeakRecordingPool::with_drift_tracker(tracker, label). The duplicate impls and tests and the unused tracker() accessor are gone.

[comphead]

self.log_threshold as isize wraps for values above isize::MAX, so with_log_threshold(usize::MAX) or --memory-drift-log-threshold 18446744073709551615 logs on every reservation change instead of never (checked with a copy of observe). saturating_add_unsigned avoids the cast:

Suggested change

	let rose = untracked >= last.saturating_add(self.log_threshold as isize);
	let rose = untracked >= last.saturating_add_unsigned(self.log_threshold);

[andygrove]

Fixed in f79c1f9 with saturating_add_unsigned. I added a test that a usize::MAX threshold never logs.

[andygrove]

@comphead I addressed the smaller points in f79c1f9:

• allocated_bytes is no longer public, and CountingAllocator no longer takes a type parameter; it always uses System.
• I merged the negative-drift test into the threshold test, and the tracker() test helper now takes a threshold.
• The Substrait round-trip engine now calls rewrap_replaced_pool too. With the call, information_schema.slt and set_variable.slt log 40 and 5 drift lines in round-trip mode. Without it, they log 37 and 2.
• The --memory-drift help text now names RUST_LOG=datafusion_execution::memory_pool=info and says that only rises of --memory-drift-log-threshold bytes are logged.

[comphead]

@andygrove please fix the doc format

[andygrove]

@comphead Fixed in 4cbebae. The MemoryPool doc links in drift.rs broke when I removed the import. cargo doc --document-private-items --no-deps --workspace with -D warnings now passes locally.