Copilot commented on code in PR #3809:
URL: https://github.com/apache/celeborn/pull/3809#discussion_r3845253285
##########
worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/TierWriter.scala:
##########
@@ -88,6 +88,32 @@ abstract class TierWriterBase(
protected def writeInternal(buf: ByteBuf): Unit
+ protected def accountAddedBytes(numBytes: Int): Unit
+
+ protected def appendToFlushBuffer(buf: ByteBuf): Unit = {
+ val numBytes = buf.readableBytes()
+ val writerIndexBefore = flushBuffer.writerIndex()
+ val numComponentsBefore = flushBuffer.numComponents()
+ buf.retain()
+ try {
+ flushBuffer.addComponent(true, buf)
+ } catch {
+ case t: Throwable =>
+ // addComponent can fail after taking ownership while consolidating
components.
+ val componentAdded =
+ flushBuffer.writerIndex() != writerIndexBefore ||
+ flushBuffer.numComponents() != numComponentsBefore
+ if (componentAdded) {
+ accountAddedBytes(numBytes)
+ }
+ logError(
+ s"Failed to add $numBytes bytes to flush buffer for shuffle
$shuffleKey file $filename.",
+ t)
+ throw t
+ }
+ accountAddedBytes(numBytes)
+ }
Review Comment:
`buf.retain()` is performed before `addComponent`, but on failures that
happen *before* the component is inserted/owned (`componentAdded == false`),
the extra retained reference is never released. This can leak a reference-count
on `buf` in the pre-insertion failure path. Fix by releasing the retained
reference when `componentAdded` is false (e.g., in the catch block or via a
finally-style cleanup), while still not releasing it when the composite
actually took ownership.
##########
worker/src/test/scala/org/apache/celeborn/service/deploy/worker/storage/TierWriterSuite.scala:
##########
@@ -37,6 +37,58 @@ import org.apache.celeborn.service.deploy.worker.WorkerSource
import org.apache.celeborn.service.deploy.worker.memory.MemoryManager
class TierWriterSuite extends AnyFunSuite with BeforeAndAfterEach {
+ private class LargeReadableByteBuf(buffer: ByteBuf, virtualSize: Int)
+ extends DuplicatedByteBuf(buffer) {
+ override def capacity(): Int = virtualSize
+
+ override def readableBytes(): Int = virtualSize
+ }
+
+ private class OomByteBufAllocator(failure: OutOfMemoryError)
+ extends AbstractByteBufAllocator(false) {
+ override protected def newHeapBuffer(initialCapacity: Int, maxCapacity:
Int): ByteBuf =
+ throw failure
+
+ override protected def newDirectBuffer(initialCapacity: Int, maxCapacity:
Int): ByteBuf =
+ throw failure
+
+ override def isDirectBufferPooled(): Boolean = false
+ }
+
+ private def capacityOverflowFlushBuffer(): CompositeByteBuf = {
+ val flushBuffer = new CompositeByteBuf(
+ UnpooledByteBufAllocator.DEFAULT,
+ false,
+ Int.MaxValue)
+ val component = UnpooledByteBufAllocator.DEFAULT.buffer(1,
Int.MaxValue).writeByte(0)
+ flushBuffer.addComponent(
+ true,
+ new LargeReadableByteBuf(component, Int.MaxValue - 512))
+ flushBuffer
+ }
+
+ private def consolidationOomFlushBuffer(failure: OutOfMemoryError):
CompositeByteBuf = {
+ val flushBuffer = new CompositeByteBuf(new OomByteBufAllocator(failure),
false, 1)
+ flushBuffer.addComponent(true, UnpooledByteBufAllocator.DEFAULT.buffer(0,
0))
+ flushBuffer
+ }
+
+ private def restoreCounters(memoryCounterBefore: Long, diskCounterBefore:
Long): Unit = {
+ val memoryManager = MemoryManager.instance()
+ val memoryCounterDelta = memoryManager.getMemoryFileStorageCounter -
memoryCounterBefore
+ if (memoryCounterDelta > 0) {
+ memoryManager.releaseMemoryFileStorage(memoryCounterDelta.toInt)
+ } else if (memoryCounterDelta < 0) {
+ memoryManager.incrementMemoryFileStorage((-memoryCounterDelta).toInt)
+ }
+ val diskCounterDelta = memoryManager.getDiskBufferCounter.get() -
diskCounterBefore
+ if (diskCounterDelta > 0) {
+ memoryManager.releaseDiskBuffer(diskCounterDelta.toInt)
+ } else if (diskCounterDelta < 0) {
+ memoryManager.incrementDiskBuffer((-diskCounterDelta).toInt)
Review Comment:
`memoryCounterDelta.toInt` / `diskCounterDelta.toInt` will silently overflow
if the delta exceeds `Int.MaxValue`, which can make cleanup incorrect and hide
test pollution. Consider guarding with a range check (fail fast in tests) or
releasing/incrementing in bounded chunks to keep the cleanup logic correct even
if a future test causes a large delta.
##########
worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/TierWriter.scala:
##########
@@ -88,6 +88,32 @@ abstract class TierWriterBase(
protected def writeInternal(buf: ByteBuf): Unit
+ protected def accountAddedBytes(numBytes: Int): Unit
+
+ protected def appendToFlushBuffer(buf: ByteBuf): Unit = {
+ val numBytes = buf.readableBytes()
+ val writerIndexBefore = flushBuffer.writerIndex()
+ val numComponentsBefore = flushBuffer.numComponents()
+ buf.retain()
+ try {
+ flushBuffer.addComponent(true, buf)
+ } catch {
+ case t: Throwable =>
+ // addComponent can fail after taking ownership while consolidating
components.
+ val componentAdded =
+ flushBuffer.writerIndex() != writerIndexBefore ||
+ flushBuffer.numComponents() != numComponentsBefore
+ if (componentAdded) {
+ accountAddedBytes(numBytes)
+ }
+ logError(
+ s"Failed to add $numBytes bytes to flush buffer for shuffle
$shuffleKey file $filename.",
+ t)
+ throw t
+ }
Review Comment:
Catching `Throwable` here will also run this logic for fatal errors (e.g.,
`InterruptedException`, `VirtualMachineError` subclasses aside from the
intended OOM/consolidation cases). It’s safer to catch
`scala.util.control.NonFatal` (and optionally handle `OutOfMemoryError`
explicitly) to avoid doing extra work (accounting + logging) during fatal
conditions.
--
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]
