Jasraj-Dhanoa opened a new pull request, #59177: URL: https://github.com/apache/spark/pull/59177
### What changes were proposed in this pull request? `StreamingShuffleWriter` can leak an off-heap buffer when a write fails partway through. This PR releases the buffer on three error paths where it is referenced only by a local variable, or where the code that releases it is skipped: 1. **Per-record serialization in `write()`.** The buffer has been taken out of its shard but not yet parked back or handed to `send()`. If the serializer throws while writing the key/value (or while computing the checksum), the buffer is now released and its `BUFFER_SIZE` permits are returned to `allocatedBufferBytesSemaphore` before the exception is rethrown. 2. **Pre-send processing in `ShardState.send()`.** Closing the serialization stream, reading the data size and computing the checksum all run before the buffer is handed to the async `send()`, whose completion callback owns the release/pool logic. If any of these throw, the buffer is now released and its permits are returned before rethrowing. 3. **Encoding in `ShardState.send(message, done)`.** Allocating the outgoing frame and encoding the message into it run before the frame is handed to the client. If either throws (for example, an `OutOfDirectMemoryError` from Netty's pooled allocator), the frame and the message were released but the completion callback `done()` was never run, so the data buffer and its permits leaked. `done()` is now run on this path too, after the message is released, matching the other failure paths in this method. The hand-off to the async `send()` is intentionally left outside the first two guards: once `send()` is called it owns the buffer, and releasing it again would corrupt the refcount. ### Why are the changes needed? The task-completion listener (`cleanupResources()`) only reclaims buffers that are pooled (`bufferPool`) or parked in a shard (`ShardState.cancel()`). A buffer that is referenced only by a local variable when an exception is thrown is invisible to it. These buffers are allocated by Netty without a cleaner, so they are not reclaimed by GC either. A leaked buffer stays lost until the executor process restarts. The executor itself survives such a task failure (including an `OutOfDirectMemoryError`) and keeps running tasks, so each failed write permanently reduces the off-heap memory available to later tasks on that executor. The file is byte-for-byte identical on `branch-4.3`, so this may be worth backporting before 4.3.0 is released. The bug was reproduced on master only. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Added four tests to `StreamingShuffleSuite`. Three use a test serializer that captures the buffer it writes into and fails at a chosen point: - `writer releases the in-flight buffer when serialization fails mid-write`: the serializer throws on the first value. Asserts the captured buffer's `refCnt` drops to 0 and all semaphore permits are returned. - `writer releases the in-flight buffer when the stream close fails on send`: the serializer throws on `close()` inside `ShardState.send()`. Asserts the captured buffer's `refCnt` drops to 0. - `shard send runs the completion callback when encoding fails`: calls `ShardState.send(message, done)` with a `DataMessage` whose data buffer throws while being encoded. Asserts `done()` runs exactly once and the message's reference is released. - `executor stays healthy after a task fails with an OutOfMemoryError`: the serializer throws an `OutOfMemoryError` mid-write. Asserts the buffer is released, then runs a new shuffle write and read on the same process and checks the data. The three leak tests fail without their fix (the buffer's `refCnt` stays at 1, one buffer's worth of permits is missing, or `done()` is never run) and pass with it. All of `StreamingShuffleSuite` passes. ``` build/sbt 'core/testOnly org.apache.spark.shuffle.streaming.StreamingShuffleSuite' ``` ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Claude Opus) This pull request and its description were written by Isaac. -- 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]
