DiskManager disk-usage counter does not return to zero after a `max_temp_directory_size` error

[fwojciec]

Describe the bug

DiskManager counts the total size of all spill files in one shared counter (used_disk_space). Each RefCountedTempFile records its own size. When a file drops, Drop subtracts the recorded size from the shared counter.

In 54.1, RefCountedTempFile::update_disk_usage does its steps in this order:

1. It charges the shared counter for the new file size.
2. It returns an error if the counter is above max_temp_directory_size.
3. It records the new size on the file.

When step 2 returns the error, step 3 does not run. The file keeps its old recorded size. Drop then subtracts too little. The difference stays in the shared counter until the process stops.

The effect grows over time. After one quota error, the manager rejects spills that fit on disk. Each rejected spill charges more bytes that are never released. On a long-lived process, spill capacity only decreases. Only a restart resets the counter.

To Reproduce

Standalone reproduction: https://github.com/fwojciec/datafusion-disk-usage-leak

cargo run uses only datafusion-execution 54.1.0 and prints:

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: ...
[BUG] all files dropped, temp dir empty, counter = 600 (expected 0)
[BUG] 500-byte spill into an empty dir rejected: Resources exhausted: ...
[BUG] counter after the rejected attempt = 1100 (leak compounds; only a process restart resets it)

cargo run --example spill_e2e --features e2e shows the same leak through the public spill path (SpillManager::spill_record_batch_and_finish): 128 bytes remain on the counter after all spill files drop.

Expected behavior

After a quota error, the counter returns to the true disk usage when the files drop. A rejected spill does not decrease future spill capacity.

Additional context

• Affected: released 54.x (datafusion-execution 54.1.0 verified). The fix there is small: record the new file size before the quota check. We run 54.1 with that one-statement reorder, and the reproduction then prints counter = 0 at every step.
• On main, feat: introduce pluggable SpillFile trait and TempFileFactory for custom spill backends #21882 removed update_disk_usage. The quota path of the new FileSpillWriter::write rolls the counter back correctly. But its write_all error path has the same defect: it charges the global counter before the write and does not release it when the write fails, and the per-file counter is not updated, so Drop subtracts too little. I can open a PR for that path if it is useful.

AI disclosure

Claude (Fable 5) found this bug during work on our DataFusion-based service, and it wrote the reproduction and this report. I reviewed the analysis and the reproduction code, ran the reproduction against stock 54.1.0, and confirmed the output. Per your AI-assisted contributions policy: I understand the accounting defect end-to-end; I have not yet studied the main refactor in the same depth.

[Himanshu-2005-code]

I'd like to work on this issue.

I'll investigate the FileSpillWriter::write_all error path on main, add a regression test reproducing the disk-usage accounting leak, and submit a fix.