DataFusion 54.1: DiskManager disk-usage counter leaks after a quota error

Minimal reproduction of an accounting leak in datafusion-execution 54.1.0.

The bug

DiskManager tracks the total size of all spill files in one shared counter (used_disk_space). Each RefCountedTempFile also records its own size. When a file is dropped, its recorded size is subtracted from the shared counter.

RefCountedTempFile::update_disk_usage performs its steps in this order:

1. Charge the shared counter for the file's growth (subtract the old recorded size, add the new actual size).
2. If the shared counter now exceeds max_temp_directory_size, return an error.
3. Record the new size on the file.

When step 2 returns the error, step 3 does not run. The shared counter was charged for the new size, but the file still records the old one. When the file is dropped, only the old size is subtracted. The difference stays in the shared counter for the life of the process.

Why it matters

The leak is a ratchet:

• After one genuine quota trip, later spills are rejected even when the temp directory is empty.
• Each rejected attempt charges its own bytes to the counter and never releases them, so the leak compounds.
• A long-lived process with max_temp_directory_size set loses spill capacity monotonically. Only a restart resets the counter.

Run the reproduction

cargo run

Observed output against crates.io datafusion-execution 54.1.0:

cap = 1000 bytes; counter at start = 0
[ok]  wrote 400 bytes, counter = 400
[ok]  dropped the file, counter = 0 (expected 0)
[ok]  second write tripped the cap: Resources exhausted: The used disk space during the spilling process has exceeded the allowable limit of 1000.0 B. Please try increasing the config: `datafusion.runtime.max_temp_directory_size`.
[BUG] all files dropped, temp dir empty, counter = 600 (expected 0)
[BUG] 500-byte spill into an empty dir rejected: Resources exhausted: The used disk space during the spilling process has exceeded the allowable limit of 1000.0 B. Please try increasing the config: `datafusion.runtime.max_temp_directory_size`.
[BUG] counter after the rejected attempt = 1100 (leak compounds; only a process restart resets it)

The same leak is reachable end-to-end through the public spill path (SpillManager::spill_record_batch_and_finish); that variant needs the full datafusion crate, so it is behind a feature:

cargo run --example spill_e2e --features e2e

The fix

Reorder the steps: record the file's new size (step 3) before the quota check (step 2). Then Drop subtracts exactly what was charged, and the counter returns to baseline after an error. We run 54.1 with this one-statement reorder applied and the reproduction above prints counter = 0 at every step.

Status on main

#21882 rewrote this code (update_disk_usage is gone), so the next release does not contain the path shown here. The rewritten FileSpillWriter::write rolls the counter back correctly on the quota path, but its write_all IO-error path still leaks the same way: the global counter is charged before the write, and neither the rollback nor the per-file record happens on that error.