Kalvin2077 commented on code in PR #3809:
URL: https://github.com/apache/celeborn/pull/3809#discussion_r3850952507
##########
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:
Fixed.
--
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]
