LuciferYang opened a new pull request, #56920:
URL: https://github.com/apache/spark/pull/56920
### What changes were proposed in this pull request?
`OneForOneBlockFetcher.DownloadCallback.onComplete` handles a chunk that has
been
streamed to a temporary file on disk. It calls `channel.closeAndRead()` to
close the
download channel and obtain a `ManagedBuffer` over the file, hands that
buffer to the
listener, and then either registers the temp file for later cleanup or
deletes it:
```java
public void onComplete(String streamId) throws IOException {
listener.onBlockFetchSuccess(blockIds[chunkIndex], channel.closeAndRead());
if (!downloadFileManager.registerTempFileToClean(targetFile)) {
targetFile.delete();
}
}
```
If `channel.closeAndRead()` throws an `IOException`, `onComplete` rethrew
without any
cleanup, so the temporary file was left on disk (it is registered for
cleanup only
*after* a successful read) and the channel could be left open. The sibling
`onFailure`
callback already performs this cleanup (`channel.close()` then
`targetFile.delete()`);
the success path did not.
This PR guards `closeAndRead()` so that, on an `IOException`, the temp file
is deleted
and the channel is closed (in case it was left open) before rethrowing,
mirroring
`onFailure`. The buffer is handed to the listener only after
`closeAndRead()` succeeds,
so on the happy path ownership of the file/buffer still transfers to the
listener
exactly once and behavior is unchanged.
```java
public void onComplete(String streamId) throws IOException {
ManagedBuffer buffer;
try {
buffer = channel.closeAndRead();
} catch (IOException e) {
// closeAndRead() failed (typically from the channel close) before the
buffer was handed
// off, so the temp file was never registered for cleanup. Delete it,
and close the
// channel in case it was left open, mirroring onFailure, instead of
leaking them.
channel.close();
targetFile.delete();
throw e;
}
listener.onBlockFetchSuccess(blockIds[chunkIndex], buffer);
if (!downloadFileManager.registerTempFileToClean(targetFile)) {
targetFile.delete();
}
}
```
Both in-tree `DownloadFileWritableChannel` implementations
(`SimpleDownloadFile.SimpleDownloadWritableChannel` and
`BlockManager.EncryptedDownloadWritableChannel`) implement `closeAndRead()`
as
"close the channel, then construct the buffer", so the most realistic source
of an
`IOException` is the channel close itself; the added `channel.close()` is
idempotent
and harmless when the channel is already closed.
### Why are the changes needed?
When fetching shuffle blocks to disk, an `IOException` from `closeAndRead()`
orphaned
the temporary block file (and potentially leaked an open channel/file
descriptor). The
external shuffle service and executors are long-lived, so these orphaned
temp files
accumulate under the local directories until the process exits. This aligns
the success
path with the already-correct cleanup in `onFailure`.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Added a unit test
`testFetchToDiskClosesAndDeletesTempFileWhenCloseAndReadFails` in
`OneForOneBlockFetcherSuite` that drives the fetch-to-disk path with a
`DownloadFileWritableChannel` whose `closeAndRead()` throws, and verifies
that the
channel is closed, the temp file is deleted, the file is never registered
for later
cleanup, and the buffer is never handed to the listener. The test fails
without the
fix (the channel is never closed) and passes with it.
`build/sbt 'network-shuffle/testOnly *OneForOneBlockFetcherSuite'` -> 14
tests pass.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
