This is an automated email from the ASF dual-hosted git repository.
ethanfeng pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new 5654c62f3 [CELEBORN-347][Flink] fix memory leak and refactor
BufferStreamManager (#1282)
5654c62f3 is described below
commit 5654c62f35b45bac1b564a31e790a5c43122b0b8
Author: Shuang <[email protected]>
AuthorDate: Tue Feb 28 15:18:59 2023 +0800
[CELEBORN-347][Flink] fix memory leak and refactor BufferStreamManager
(#1282)
---
.../common/network/server/BufferStreamManager.java | 571 ++++++---------------
.../common/network/server/DataPartitionReader.java | 448 ++++++++++++++++
.../network/server/memory/BufferRecycler.java | 44 ++
.../server/memory/ReadBufferDispatcher.java | 11 +-
.../common/network/server/memory/Recycler.java | 30 ++
.../network/server/memory/WrappedDataBuffer.java | 40 ++
6 files changed, 721 insertions(+), 423 deletions(-)
diff --git
a/common/src/main/java/org/apache/celeborn/common/network/server/BufferStreamManager.java
b/common/src/main/java/org/apache/celeborn/common/network/server/BufferStreamManager.java
index e48f28765..9e367776d 100644
---
a/common/src/main/java/org/apache/celeborn/common/network/server/BufferStreamManager.java
+++
b/common/src/main/java/org/apache/celeborn/common/network/server/BufferStreamManager.java
@@ -20,9 +20,7 @@ package org.apache.celeborn.common.network.server;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
-import java.nio.ByteBuffer;
import java.nio.channels.FileChannel;
-import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
@@ -32,13 +30,14 @@ import java.util.PriorityQueue;
import java.util.Queue;
import java.util.Random;
import java.util.Set;
+import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.DelayQueue;
+import java.util.concurrent.Delayed;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingDeque;
-import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Consumer;
@@ -47,22 +46,18 @@ import javax.annotation.concurrent.GuardedBy;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import io.netty.buffer.ByteBuf;
import io.netty.channel.Channel;
-import io.netty.channel.ChannelFutureListener;
import org.apache.commons.io.IOUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.meta.FileInfo;
-import org.apache.celeborn.common.network.protocol.BacklogAnnouncement;
-import org.apache.celeborn.common.network.protocol.ReadData;
+import org.apache.celeborn.common.network.server.memory.BufferRecycler;
import org.apache.celeborn.common.network.server.memory.MemoryManager;
-import org.apache.celeborn.common.util.Utils;
public class BufferStreamManager {
private static final Logger logger =
LoggerFactory.getLogger(BufferStreamManager.class);
private final AtomicLong nextStreamId;
protected final ConcurrentHashMap<Long, StreamState> streams;
- protected final ConcurrentHashMap<Long, AtomicInteger> streamCredits;
protected final ConcurrentHashMap<Long, MapDataPartition> servingStreams;
protected final ConcurrentHashMap<FileInfo, MapDataPartition>
activeMapPartitions;
protected final MemoryManager memoryManager = MemoryManager.instance();
@@ -70,10 +65,10 @@ public class BufferStreamManager {
protected int minReadBuffers;
protected int maxReadBuffers;
protected int threadsPerMountPoint;
- private final LinkedBlockingQueue<Long> recycleStreamIds = new
LinkedBlockingQueue<>();
+ private final BlockingQueue<DelayedStreamId> recycleStreamIds = new
DelayQueue<>();
@GuardedBy("lock")
- private Thread recycleThread;
+ private volatile Thread recycleThread;
private final Object lock = new Object();
@@ -98,7 +93,6 @@ public class BufferStreamManager {
public BufferStreamManager(int minReadBuffers, int maxReadBuffers, int
threadsPerMountpoint) {
nextStreamId = new AtomicLong((long) new
Random().nextInt(Integer.MAX_VALUE) * 1000);
streams = new ConcurrentHashMap<>();
- streamCredits = new ConcurrentHashMap<>();
servingStreams = new ConcurrentHashMap<>();
activeMapPartitions = new ConcurrentHashMap<>();
this.minReadBuffers = minReadBuffers;
@@ -116,7 +110,7 @@ public class BufferStreamManager {
throws IOException {
long streamId = nextStreamId.getAndIncrement();
streams.put(streamId, new StreamState(channel, fileInfo.getBufferSize()));
-
+ logger.debug("Register stream start streamId: {}, fileInfo: {}", streamId,
fileInfo);
synchronized (activeMapPartitions) {
MapDataPartition mapDataPartition = activeMapPartitions.get(fileInfo);
if (mapDataPartition == null) {
@@ -128,7 +122,7 @@ public class BufferStreamManager {
servingStreams.put(streamId, mapDataPartition);
// response streamId to channel first
callback.accept(streamId);
- mapDataPartition.setupDataPartitionReader(startSubIndex, endSubIndex,
streamId);
+ mapDataPartition.setupDataPartitionReader(startSubIndex, endSubIndex,
streamId, channel);
}
logger.debug("Register stream streamId: {}, fileInfo: {}", streamId,
fileInfo);
@@ -138,25 +132,13 @@ public class BufferStreamManager {
public void addCredit(int numCredit, long streamId) {
logger.debug("streamId: {}, add credit: {}", streamId, numCredit);
- streamCredits.compute(
- streamId,
- (aLong, atomicInteger) -> {
- if (atomicInteger == null) {
- return new AtomicInteger(numCredit);
- } else {
- atomicInteger.getAndAdd(numCredit);
- return atomicInteger;
- }
- });
-
- MapDataPartition mapDataPartition = servingStreams.get(streamId);
- if (mapDataPartition != null) {
- DataPartitionReader streamReader =
mapDataPartition.getStreamReader(streamId);
- if (streamReader != null) {
- storageFetcherPool
- .getExecutorPool(streamReader.fileInfo.getMountPoint())
- .submit(() -> streamReader.sendData());
+ try {
+ MapDataPartition mapDataPartition = servingStreams.get(streamId);
+ if (mapDataPartition != null) {
+ mapDataPartition.addReaderCredit(numCredit, streamId);
}
+ } catch (Throwable e) {
+ logger.error("streamId: {}, add credit end: {}", streamId, numCredit);
}
}
@@ -170,7 +152,7 @@ public class BufferStreamManager {
}
public void recycleStream(long streamId) {
- recycleStreamIds.add(streamId);
+ recycleStreamIds.add(new DelayedStreamId(streamId));
startRecycleThread(); // lazy start thread
}
@@ -183,13 +165,10 @@ public class BufferStreamManager {
() -> {
while (true) {
try {
- Long streamIds = recycleStreamIds.poll(100,
TimeUnit.MILLISECONDS);
- if (streamIds != null) {
- cleanResource(streamIds);
- }
- } catch (InterruptedException e) {
+ DelayedStreamId delayedStreamId =
recycleStreamIds.take();
+ cleanResource(delayedStreamId.streamId);
+ } catch (Throwable e) {
logger.warn(e.getMessage(), e);
- throw new RuntimeException(e);
}
}
},
@@ -204,25 +183,69 @@ public class BufferStreamManager {
}
public void cleanResource(Long streamId) {
- logger.debug("clean stream: {}", streamId);
- streams.remove(streamId);
- streamCredits.remove(streamId);
- MapDataPartition mapDataPartition = servingStreams.remove(streamId);
- if (mapDataPartition != null) {
- FileInfo fileInfo = mapDataPartition.fileInfo;
- mapDataPartition.removeStream(streamId);
- synchronized (activeMapPartitions) {
- if (mapDataPartition.activeStreamIds.isEmpty()) {
- for (ByteBuf buffer : mapDataPartition.buffers) {
- memoryManager.recycleReadBuffer(buffer);
+ logger.debug("received clean stream: {}", streamId);
+ if (streams.containsKey(streamId)) {
+ MapDataPartition mapDataPartition = servingStreams.get(streamId);
+ if (mapDataPartition != null) {
+ if (mapDataPartition.releaseStream(streamId)) {
+ synchronized (activeMapPartitions) {
+ if (mapDataPartition.activeStreamIds.isEmpty()) {
+ mapDataPartition.close();
+ FileInfo fileInfo = mapDataPartition.fileInfo;
+ activeMapPartitions.remove(fileInfo);
+ }
}
- activeMapPartitions.remove(fileInfo);
- mapDataPartition.close();
+ } else {
+ logger.debug("retry clean stream: {}", streamId);
+ recycleStreamIds.add(new DelayedStreamId(streamId));
}
}
}
}
+ public static class DelayedStreamId implements Delayed {
+ private static final long delayTime = 100; // 100ms
+ private long createMillis = System.currentTimeMillis();
+
+ private long streamId;
+
+ public DelayedStreamId(long streamId) {
+ this.createMillis = createMillis + delayTime;
+ this.streamId = streamId;
+ }
+
+ @Override
+ public long getDelay(TimeUnit unit) {
+ long diff = createMillis - System.currentTimeMillis();
+ return unit.convert(diff, TimeUnit.MILLISECONDS);
+ }
+
+ public long getCreateMillis() {
+ return createMillis;
+ }
+
+ @Override
+ public int compareTo(Delayed o) {
+ long otherCreateMillis = ((DelayedStreamId) o).getCreateMillis();
+ if (this.createMillis < otherCreateMillis) {
+ return -1;
+ } else if (this.createMillis > otherCreateMillis) {
+ return 1;
+ }
+
+ return 0;
+ }
+
+ @Override
+ public String toString() {
+ final StringBuilder sb = new StringBuilder("DelayedStreamId{");
+ sb.append("createMillis=").append(createMillis);
+ sb.append(", streamId=").append(streamId);
+ sb.append('}');
+ return sb.toString();
+ }
+ }
+
// this means active data partition
protected class MapDataPartition {
private final List<Long> activeStreamIds = new ArrayList<>();
@@ -233,10 +256,12 @@ public class BufferStreamManager {
new ConcurrentHashMap<>();
/** All available buffers can be used by the partition readers for
reading. */
- protected Queue<ByteBuf> buffers;
+ private Queue<ByteBuf> buffers;
- protected FileChannel dataFileChanel;
- protected FileChannel indexChannel;
+ private FileChannel dataFileChanel;
+ private FileChannel indexChannel;
+
+ private boolean isReleased;
public MapDataPartition(FileInfo fileInfo) throws FileNotFoundException {
this.fileInfo = fileInfo;
@@ -246,9 +271,15 @@ public class BufferStreamManager {
}
public synchronized void setupDataPartitionReader(
- int startSubIndex, int endSubIndex, long streamId) throws IOException {
+ int startSubIndex, int endSubIndex, long streamId, Channel channel)
throws IOException {
DataPartitionReader dataPartitionReader =
- new DataPartitionReader(startSubIndex, endSubIndex, fileInfo,
streamId);
+ new DataPartitionReader(
+ startSubIndex,
+ endSubIndex,
+ fileInfo,
+ streamId,
+ channel,
+ () -> recycleStream(streamId));
dataPartitionReader.open(dataFileChanel, indexChannel);
// allocate resources when the first reader is registered
boolean allocateResources = readers.isEmpty();
@@ -256,7 +287,7 @@ public class BufferStreamManager {
streamReaders.put(streamId, dataPartitionReader);
// create initial buffers for read
- if (allocateResources) {
+ if (allocateResources && buffers == null) {
memoryManager.requestReadBuffers(
minReadBuffers,
maxReadBuffers,
@@ -268,10 +299,6 @@ public class BufferStreamManager {
}
}
- public DataPartitionReader getStreamReader(long streamId) {
- return streamReaders.get(streamId);
- }
-
// Read logic is executed on another thread.
public void onBuffer(Queue<ByteBuf> buffers) {
this.buffers = buffers;
@@ -285,393 +312,99 @@ public class BufferStreamManager {
}
public synchronized void readBuffers() {
- PriorityQueue<DataPartitionReader> sortedReaders = new
PriorityQueue<>(readers);
- while (buffers.size() > 0 && !sortedReaders.isEmpty()) {
- DataPartitionReader reader = sortedReaders.poll();
- try {
- if (!reader.readAndSend(buffers, (buffer) -> this.recycle(buffer,
buffers))) {
+ if (isReleased) {
+ // some read executor task may already be submitted to the threadpool
+ return;
+ }
+
+ try {
+ PriorityQueue<DataPartitionReader> sortedReaders = new
PriorityQueue<>(readers);
+ while (buffers != null && buffers.size() > 0 &&
!sortedReaders.isEmpty()) {
+ BufferRecycler bufferRecycler =
+ new BufferRecycler(memoryManager, (buffer) ->
this.recycle(buffer, buffers));
+ DataPartitionReader reader = sortedReaders.poll();
+ try {
+ if (!reader.readAndSend(buffers, bufferRecycler)) {
+ readers.remove(reader);
+ }
+ } catch (Throwable e) {
+ logger.error("reader exception, reader: {}, message: {}", reader,
e.getMessage(), e);
readers.remove(reader);
+ reader.recycleOnError(e);
}
- } catch (IOException e) {
- logger.error("Read thread error occurred, {}", e);
- throw new RuntimeException(e);
}
- }
- }
-
- public void triggerRead() {
- readExecutor.submit(
- () -> {
- // Key for IO schedule.
- readBuffers();
- });
- }
-
- public synchronized void addStream(long streamId) {
- activeStreamIds.add(streamId);
- }
-
- public synchronized void removeStream(long streamId) {
- activeStreamIds.remove(streamId);
- }
-
- public void close() {
- IOUtils.closeQuietly(dataFileChanel);
- IOUtils.closeQuietly(indexChannel);
- }
- }
-
- private class Buffer {
- private ByteBuf byteBuf;
- private Consumer<ByteBuf> byteBufferConsumer;
-
- public Buffer(ByteBuf byteBuf, Consumer<ByteBuf> byteBufferConsumer) {
- this.byteBuf = byteBuf;
- this.byteBufferConsumer = byteBufferConsumer;
- }
- }
-
- // this is a specific partition reader
- protected class DataPartitionReader implements
Comparable<DataPartitionReader> {
- private final ByteBuffer indexBuffer;
- private final ByteBuffer headerBuffer;
- private final int startPartitionIndex;
- private final int endPartitionIndex;
- private int numRegions;
- private int numRemainingPartitions;
- private int currentDataRegion = -1;
- private long dataConsumingOffset;
- private volatile long currentPartitionRemainingBytes;
- private boolean isClosed;
- private FileInfo fileInfo;
- private int INDEX_ENTRY_SIZE = 16;
- private long streamId;
- protected final Object lock = new Object();
-
- @GuardedBy("lock")
- protected final Queue<Buffer> buffersRead = new ArrayDeque<>();
-
- /** Whether all the data has been successfully read or not. */
- @GuardedBy("lock")
- protected boolean isFinished;
-
- /** Whether this partition reader has been released or not. */
- @GuardedBy("lock")
- protected boolean isReleased;
-
- /** Exception causing the release of this partition reader. */
- @GuardedBy("lock")
- protected Throwable errorCause;
-
- /** Whether there is any error at the consumer side or not. */
- @GuardedBy("lock")
- protected boolean isError;
-
- private FileChannel dataFileChannel;
- private FileChannel indexFileChannel;
-
- public DataPartitionReader(
- int startPartitionIndex, int endPartitionIndex, FileInfo fileInfo,
long streamId) {
- this.startPartitionIndex = startPartitionIndex;
- this.endPartitionIndex = endPartitionIndex;
-
- int indexBufferSize = 16 * (endPartitionIndex - startPartitionIndex + 1);
- this.indexBuffer = ByteBuffer.allocateDirect(indexBufferSize);
-
- this.headerBuffer = ByteBuffer.allocateDirect(16);
- this.streamId = streamId;
-
- this.fileInfo = fileInfo;
- this.isClosed = false;
- }
-
- public void open(FileChannel dataFileChannel, FileChannel
indexFileChannel) throws IOException {
- this.dataFileChannel = dataFileChannel;
- this.indexFileChannel = indexFileChannel;
- long indexFileSize = indexFileChannel.size();
- // index is (offset,length)
- long indexRegionSize = fileInfo.getNumReducerPartitions() * (long)
INDEX_ENTRY_SIZE;
- this.numRegions = Utils.checkedDownCast(indexFileSize / indexRegionSize);
-
- updateConsumingOffset();
- }
-
- public long getIndexRegionSize() {
- return fileInfo.getNumReducerPartitions() * (long) INDEX_ENTRY_SIZE;
- }
+ } catch (Throwable e) {
+ logger.error("Fatal: failed to read partition data. {}",
e.getMessage(), e);
+ for (DataPartitionReader reader : readers) {
+ reader.recycleOnError(e);
+ }
- private void readHeaderOrIndexBuffer(FileChannel channel, ByteBuffer
buffer, int length)
- throws IOException {
- Utils.checkFileIntegrity(channel, length);
- buffer.clear();
- buffer.limit(length);
- while (buffer.hasRemaining()) {
- channel.read(buffer);
+ readers.clear();
}
- buffer.flip();
}
- private void readBufferIntoReadBuffer(FileChannel channel, ByteBuf buf,
int length)
- throws IOException {
- Utils.checkFileIntegrity(channel, length);
- ByteBuffer tmpBuffer = ByteBuffer.allocate(length);
- while (tmpBuffer.hasRemaining()) {
- channel.read(tmpBuffer);
+ // for one reader only the associated channel can access
+ public void addReaderCredit(int numCredit, long streamId) {
+ DataPartitionReader streamReader = this.getStreamReader(streamId);
+ if (streamReader != null) {
+ boolean canSendWithCredit = streamReader.sendWithCredit(numCredit);
+ if (canSendWithCredit) {
+ readExecutor.submit(() -> streamReader.sendData());
+ }
}
- tmpBuffer.flip();
- buf.writeBytes(tmpBuffer);
}
- private int readBuffer(
- String filename, FileChannel channel, ByteBuffer header, ByteBuf
buffer, int headerSize)
- throws IOException {
- readHeaderOrIndexBuffer(channel, header, headerSize);
- // header is combined of mapId(4),attemptId(4),nextBatchId(4) and
totcalCompresszedLength(4)
- // we need size here,so we read length directly
- int bufferLenght = header.getInt(12);
- if (bufferLenght <= 0 || bufferLenght > buffer.capacity()) {
- logger.error("Incorrect buffer header, buffer length: {}.",
bufferLenght);
- throw new RuntimeException("File " + filename + " is corrupted");
- }
- buffer.writeBytes(header);
- readBufferIntoReadBuffer(channel, buffer, bufferLenght);
- return bufferLenght + headerSize;
+ public void triggerRead() {
+ // Key for IO schedule.
+ readExecutor.submit(() -> readBuffers());
}
- private void updateConsumingOffset() throws IOException {
- while (currentPartitionRemainingBytes == 0
- && (currentDataRegion < numRegions - 1 || numRemainingPartitions >
0)) {
- if (numRemainingPartitions <= 0) {
- ++currentDataRegion;
- numRemainingPartitions = endPartitionIndex - startPartitionIndex + 1;
-
- // read the target index entry to the target index buffer
- indexFileChannel.position(
- currentDataRegion * getIndexRegionSize()
- + (long) startPartitionIndex * INDEX_ENTRY_SIZE);
- readHeaderOrIndexBuffer(indexFileChannel, indexBuffer,
indexBuffer.capacity());
- }
-
- // get the data file offset and the data size
- dataConsumingOffset = indexBuffer.getLong();
- currentPartitionRemainingBytes = indexBuffer.getLong();
- --numRemainingPartitions;
-
- logger.debug(
- "readBuffer updateConsumingOffset, {}, {}, {}, {}",
- streamId,
- dataFileChannel.size(),
- dataConsumingOffset,
- currentPartitionRemainingBytes);
-
- // if these checks fail, the partition file must be corrupted
- if (dataConsumingOffset < 0
- || dataConsumingOffset + currentPartitionRemainingBytes >
dataFileChannel.size()
- || currentPartitionRemainingBytes < 0) {
- throw new RuntimeException("File " + fileInfo.getFilePath() + " is
corrupted");
- }
+ public void addStream(Long streamId) {
+ synchronized (activeStreamIds) {
+ activeStreamIds.add(streamId);
}
}
- public boolean readBuffer(ByteBuf buffer) throws IOException {
- try {
- dataFileChannel.position(dataConsumingOffset);
-
- int readSize =
- readBuffer(
- fileInfo.getFilePath(),
- dataFileChannel,
- headerBuffer,
- buffer,
- headerBuffer.capacity());
- currentPartitionRemainingBytes -= readSize;
-
- logger.debug(
- "readBuffer data: {}, {}, {}, {}, {}, {}",
- streamId,
- currentPartitionRemainingBytes,
- readSize,
- dataConsumingOffset,
- fileInfo.getFilePath(),
- System.identityHashCode(buffer));
-
- // if this check fails, the partition file must be corrupted
- if (currentPartitionRemainingBytes < 0) {
- throw new RuntimeException("File is corrupted");
- } else if (currentPartitionRemainingBytes == 0) {
- logger.debug(
- "readBuffer end, {}, {}, {}, {}",
- streamId,
- dataFileChannel.size(),
- dataConsumingOffset,
- currentPartitionRemainingBytes);
- int prevDataRegion = currentDataRegion;
- updateConsumingOffset();
- return prevDataRegion == currentDataRegion &&
currentPartitionRemainingBytes > 0;
- }
-
- dataConsumingOffset = dataFileChannel.position();
-
- logger.debug(
- "readBuffer run: {}, {}, {}, {}",
- streamId,
- dataFileChannel.size(),
- dataConsumingOffset,
- currentPartitionRemainingBytes);
- return true;
- } catch (Throwable throwable) {
- logger.debug("Failed to read partition file.", throwable);
- isReleased = true;
- throw throwable;
+ public void removeStream(Long streamId) {
+ synchronized (activeStreamIds) {
+ activeStreamIds.remove(streamId);
+ streamReaders.remove(streamId);
}
}
- public boolean hasRemaining() {
- return currentPartitionRemainingBytes > 0;
+ public DataPartitionReader getStreamReader(long streamId) {
+ return streamReaders.get(streamId);
}
- public synchronized boolean readAndSend(Queue<ByteBuf> bufferQueue,
Consumer<ByteBuf> consumer)
- throws IOException {
- boolean hasRemaing = hasRemaining();
- boolean continueReading = hasRemaing;
- int numDataBuffers = 0;
- while (continueReading) {
-
- ByteBuf buffer = bufferQueue.poll();
- // this is used for control bytebuf manually.
- if (buffer == null) {
- break;
- } else {
- buffer.retain();
- }
-
- try {
- continueReading = readBuffer(buffer);
- } catch (Throwable throwable) {
- memoryManager.recycleReadBuffer(buffer);
- throw throwable;
- }
-
- hasRemaing = hasRemaining();
- addBuffer(buffer, hasRemaing, consumer);
- ++numDataBuffers;
- }
- if (numDataBuffers > 0) {
- notifyBacklog(numDataBuffers);
- }
-
- if (!hasRemaing) {
- closeReader();
+ public boolean releaseStream(Long streamId) {
+ DataPartitionReader dataPartitionReader = streamReaders.get(streamId);
+ dataPartitionReader.release();
+ if (dataPartitionReader.isFinished()) {
+ logger.info("release all for stream: {}", streamId);
+ removeStream(streamId);
+ streams.remove(streamId);
+ servingStreams.remove(streamId);
+ return true;
}
- return hasRemaing;
+ return false;
}
- protected void addBuffer(ByteBuf buffer, boolean hasRemaining,
Consumer<ByteBuf> consumer) {
- if (buffer == null) {
- return;
- }
- final boolean recycleBuffer;
- boolean notifyDataAvailable = false;
- final Throwable throwable;
- synchronized (lock) {
- recycleBuffer = isReleased || isFinished || isError;
- throwable = errorCause;
- isFinished = !hasRemaining;
-
- if (!recycleBuffer) {
- notifyDataAvailable = buffersRead.isEmpty();
- buffersRead.add(new Buffer(buffer, consumer));
- }
- }
+ public void close() {
+ logger.info("release map data partition {}", fileInfo);
- if (recycleBuffer) {
- memoryManager.recycleReadBuffer(buffer);
- throw new RuntimeException("Partition reader has been failed or
finished.", throwable);
- }
- if (!buffersRead.isEmpty()) {
- sendData();
- }
- }
+ IOUtils.closeQuietly(dataFileChanel);
+ IOUtils.closeQuietly(indexChannel);
- public synchronized void sendData() {
- while (!buffersRead.isEmpty()) {
- logger.debug("senddata streamid:{}, {}", streamId,
streamCredits.get(streamId));
- if (streamCredits.get(streamId).get() > 0) {
- Buffer readBuf = buffersRead.poll();
- logger.debug(
- "send {} to stream {} ,fileInfo: {}",
- readBuf.byteBuf.readableBytes(),
- streamId,
- fileInfo.getFilePath());
- ReadData readData = new ReadData(streamId, buffersRead.size(), 0,
readBuf.byteBuf);
- streams
- .get(streamId)
- .associatedChannel
- .writeAndFlush(readData)
- .addListener(
- (ChannelFutureListener)
- future -> {
- if (!future.isSuccess()) {
- isError = true;
- errorCause = future.cause();
- memoryManager.recycleReadBuffer(readBuf.byteBuf);
- throw new RuntimeException(future.cause());
- } else {
- // netty has release the reference
- }
- logger.debug(
- "recycle {} , {} to stream {} ,fileinfo {}",
- readBuf.byteBuf.readableBytes(),
- System.identityHashCode(readBuf.byteBuf),
- streamId,
- fileInfo.getFilePath());
- readBuf.byteBufferConsumer.accept(readBuf.byteBuf);
- });
- logger.debug(
- "streamId: {}, decrement credit: {}",
- streamId,
- streamCredits.get(streamId).decrementAndGet());
- } else {
- logger.debug("streamId: {}, no credit: {}", streamId,
streamCredits.get(streamId));
- // no credit
- break;
+ if (this.buffers != null) {
+ for (ByteBuf buffer : this.buffers) {
+ memoryManager.recycleReadBuffer(buffer);
}
}
- logger.debug(
- "streamId: {}, remaining: {}, bufferSize: {}",
- streamId,
- currentPartitionRemainingBytes,
- buffersRead.size());
-
- if (isClosed && buffersRead.isEmpty()) {
- recycleStream(streamId);
- }
- }
-
- public void closeReader() {
- isClosed = true;
- logger.debug("Closed read for stream {}", this.streamId);
- }
-
- protected void notifyBacklog(int backlog) {
- logger.debug("stream manager streamid {} backlog:{}", streamId, backlog);
- StreamState streamState = streams.get(streamId);
- if (streamState == null) {
- logger.warn("StreamId :{} already closed", streamId);
- } else {
- streamState.associatedChannel.writeAndFlush(new
BacklogAnnouncement(streamId, backlog));
- }
- }
-
- public long getPriority() {
- return dataConsumingOffset;
- }
+ this.buffers = null;
- @Override
- public int compareTo(DataPartitionReader that) {
- return Long.compare(getPriority(), that.getPriority());
+ isReleased = true;
}
}
diff --git
a/common/src/main/java/org/apache/celeborn/common/network/server/DataPartitionReader.java
b/common/src/main/java/org/apache/celeborn/common/network/server/DataPartitionReader.java
new file mode 100644
index 000000000..612712147
--- /dev/null
+++
b/common/src/main/java/org/apache/celeborn/common/network/server/DataPartitionReader.java
@@ -0,0 +1,448 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.common.network.server;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.FileChannel;
+import java.util.ArrayDeque;
+import java.util.Queue;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import javax.annotation.concurrent.GuardedBy;
+
+import io.netty.buffer.ByteBuf;
+import io.netty.channel.Channel;
+import io.netty.channel.ChannelFutureListener;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.common.meta.FileInfo;
+import org.apache.celeborn.common.network.protocol.BacklogAnnouncement;
+import org.apache.celeborn.common.network.protocol.ReadData;
+import org.apache.celeborn.common.network.server.memory.Recycler;
+import org.apache.celeborn.common.network.server.memory.WrappedDataBuffer;
+import org.apache.celeborn.common.util.Utils;
+
+public class DataPartitionReader implements Comparable<DataPartitionReader> {
+ private static final Logger logger =
LoggerFactory.getLogger(DataPartitionReader.class);
+
+ private final ByteBuffer indexBuffer;
+ private final ByteBuffer headerBuffer;
+ private final int startPartitionIndex;
+ private final int endPartitionIndex;
+ private int numRegions;
+ private int numRemainingPartitions;
+ private int currentDataRegion = -1;
+ private long dataConsumingOffset;
+ private volatile long currentPartitionRemainingBytes;
+ private FileInfo fileInfo;
+ private int INDEX_ENTRY_SIZE = 16;
+ private long streamId;
+ protected final Object lock = new Object();
+
+ private AtomicInteger credits = new AtomicInteger();
+
+ @GuardedBy("lock")
+ protected final Queue<WrappedDataBuffer> buffersRead = new ArrayDeque<>();
+
+ /** Whether all the data has been successfully read or not. */
+ @GuardedBy("lock")
+ private boolean isClosed;
+
+ /** Whether this partition reader has been released or not. */
+ @GuardedBy("lock")
+ protected boolean isReleased;
+
+ /** Exception causing the release of this partition reader. */
+ @GuardedBy("lock")
+ protected Throwable errorCause;
+
+ /** Whether there is any error at the consumer side or not. */
+ @GuardedBy("lock")
+ protected boolean isError;
+
+ private FileChannel dataFileChannel;
+ private FileChannel indexFileChannel;
+
+ private Channel associatedChannel;
+
+ private Runnable recycleStream;
+
+ private AtomicInteger numInFlightRequests = new AtomicInteger(0);
+
+ public DataPartitionReader(
+ int startPartitionIndex,
+ int endPartitionIndex,
+ FileInfo fileInfo,
+ long streamId,
+ Channel associatedChannel,
+ Runnable recycleStream) {
+ this.startPartitionIndex = startPartitionIndex;
+ this.endPartitionIndex = endPartitionIndex;
+
+ int indexBufferSize = 16 * (endPartitionIndex - startPartitionIndex + 1);
+ this.indexBuffer = ByteBuffer.allocateDirect(indexBufferSize);
+
+ this.headerBuffer = ByteBuffer.allocateDirect(16);
+ this.streamId = streamId;
+ this.associatedChannel = associatedChannel;
+ this.recycleStream = recycleStream;
+
+ this.fileInfo = fileInfo;
+ this.isClosed = false;
+ }
+
+ public void open(FileChannel dataFileChannel, FileChannel indexFileChannel)
throws IOException {
+ this.dataFileChannel = dataFileChannel;
+ this.indexFileChannel = indexFileChannel;
+ long indexFileSize = indexFileChannel.size();
+ // index is (offset,length)
+ long indexRegionSize = fileInfo.getNumReducerPartitions() * (long)
INDEX_ENTRY_SIZE;
+ this.numRegions = Utils.checkedDownCast(indexFileSize / indexRegionSize);
+
+ updateConsumingOffset();
+ }
+
+ public boolean sendWithCredit(int credit) {
+ int oldCredit = credits.getAndAdd(credit);
+ if (oldCredit == 0) {
+ return true;
+ }
+
+ return false;
+ }
+
+ public synchronized boolean readAndSend(Queue<ByteBuf> bufferQueue, Recycler
bufferRecycler)
+ throws IOException {
+ boolean hasRemaining = hasRemaining();
+ boolean continueReading = hasRemaining;
+ int numDataBuffers = 0;
+ while (continueReading) {
+
+ ByteBuf buffer = bufferQueue.poll();
+ // this is used for control bytebuf manually.
+ if (buffer == null) {
+ break;
+ } else {
+ buffer.retain();
+ }
+
+ try {
+ continueReading = readBuffer(buffer);
+ } catch (Throwable throwable) {
+ bufferRecycler.release(buffer);
+ throw throwable;
+ }
+
+ hasRemaining = hasRemaining();
+ addBuffer(buffer, hasRemaining, bufferRecycler);
+ ++numDataBuffers;
+ }
+ if (numDataBuffers > 0) {
+ notifyBacklog(numDataBuffers);
+ }
+
+ if (!hasRemaining) {
+ closeReader();
+ }
+
+ return hasRemaining;
+ }
+
+ private void addBuffer(ByteBuf buffer, boolean hasRemaining, Recycler
bufferRecycler) {
+ if (buffer == null) {
+ return;
+ }
+ final boolean recycleBuffer;
+ boolean notifyDataAvailable = false;
+ final Throwable throwable;
+ synchronized (lock) {
+ recycleBuffer = isReleased || isClosed || isError;
+ throwable = errorCause;
+ isClosed = !hasRemaining;
+
+ if (!recycleBuffer) {
+ notifyDataAvailable = buffersRead.isEmpty();
+ buffersRead.add(new WrappedDataBuffer(buffer, bufferRecycler));
+ }
+ }
+
+ if (recycleBuffer) {
+ bufferRecycler.release(buffer);
+ throw new RuntimeException("Partition reader has been failed or
finished.", throwable);
+ }
+ if (notifyDataAvailable) {
+ sendData();
+ }
+ }
+
+ public synchronized void sendData() {
+ while (!buffersRead.isEmpty() && credits.get() > 0) {
+ WrappedDataBuffer wrappedBuffer;
+ synchronized (lock) {
+ wrappedBuffer = buffersRead.poll();
+ numInFlightRequests.incrementAndGet();
+ }
+
+ ByteBuf byteBuf = wrappedBuffer.byteBuf;
+ int backlog = buffersRead.size();
+ int readableBytes = byteBuf.readableBytes();
+ logger.debug("send data start: {}, {}, {}", streamId, readableBytes,
backlog);
+ ReadData readData = new ReadData(streamId, backlog, 0, byteBuf);
+ associatedChannel
+ .writeAndFlush(readData)
+ .addListener(
+ (ChannelFutureListener)
+ future -> {
+ try {
+ if (!future.isSuccess()) {
+ wrappedBuffer.release();
+ recycleOnError(future.cause());
+ } else {
+ wrappedBuffer.recycle();
+ }
+ } finally {
+ logger.debug("send data start: {}, {}, {}", streamId,
readableBytes, backlog);
+ numInFlightRequests.decrementAndGet();
+ }
+ });
+
+ credits.decrementAndGet();
+ }
+
+ if (isClosed && buffersRead.isEmpty()) {
+ recycle();
+ }
+ }
+
+ private long getIndexRegionSize() {
+ return fileInfo.getNumReducerPartitions() * (long) INDEX_ENTRY_SIZE;
+ }
+
+ private void readHeaderOrIndexBuffer(FileChannel channel, ByteBuffer buffer,
int length)
+ throws IOException {
+ Utils.checkFileIntegrity(channel, length);
+ buffer.clear();
+ buffer.limit(length);
+ while (buffer.hasRemaining()) {
+ channel.read(buffer);
+ }
+ buffer.flip();
+ }
+
+ private void readBufferIntoReadBuffer(FileChannel channel, ByteBuf buf, int
length)
+ throws IOException {
+ Utils.checkFileIntegrity(channel, length);
+ ByteBuffer tmpBuffer = ByteBuffer.allocate(length);
+ while (tmpBuffer.hasRemaining()) {
+ channel.read(tmpBuffer);
+ }
+ tmpBuffer.flip();
+ buf.writeBytes(tmpBuffer);
+ }
+
+ private int readBuffer(
+ String filename, FileChannel channel, ByteBuffer header, ByteBuf buffer,
int headerSize)
+ throws IOException {
+ readHeaderOrIndexBuffer(channel, header, headerSize);
+ // header is combined of mapId(4),attemptId(4),nextBatchId(4) and total
Compresszed Length(4)
+ // we need size here,so we read length directly
+ int bufferLength = header.getInt(12);
+ if (bufferLength <= 0 || bufferLength > buffer.capacity()) {
+ logger.error("Incorrect buffer header, buffer length: {}.",
bufferLength);
+ throw new RuntimeException("File " + filename + " is corrupted");
+ }
+ buffer.writeBytes(header);
+ readBufferIntoReadBuffer(channel, buffer, bufferLength);
+ return bufferLength + headerSize;
+ }
+
+ private void updateConsumingOffset() throws IOException {
+ while (currentPartitionRemainingBytes == 0
+ && (currentDataRegion < numRegions - 1 || numRemainingPartitions > 0))
{
+ if (numRemainingPartitions <= 0) {
+ ++currentDataRegion;
+ numRemainingPartitions = endPartitionIndex - startPartitionIndex + 1;
+
+ // read the target index entry to the target index buffer
+ indexFileChannel.position(
+ currentDataRegion * getIndexRegionSize()
+ + (long) startPartitionIndex * INDEX_ENTRY_SIZE);
+ readHeaderOrIndexBuffer(indexFileChannel, indexBuffer,
indexBuffer.capacity());
+ }
+
+ // get the data file offset and the data size
+ dataConsumingOffset = indexBuffer.getLong();
+ currentPartitionRemainingBytes = indexBuffer.getLong();
+ --numRemainingPartitions;
+
+ logger.debug(
+ "readBuffer updateConsumingOffset, {}, {}, {}, {}",
+ streamId,
+ dataFileChannel.size(),
+ dataConsumingOffset,
+ currentPartitionRemainingBytes);
+
+ // if these checks fail, the partition file must be corrupted
+ if (dataConsumingOffset < 0
+ || dataConsumingOffset + currentPartitionRemainingBytes >
dataFileChannel.size()
+ || currentPartitionRemainingBytes < 0) {
+ throw new RuntimeException("File " + fileInfo.getFilePath() + " is
corrupted");
+ }
+ }
+ }
+
+ private boolean readBuffer(ByteBuf buffer) throws IOException {
+ try {
+ dataFileChannel.position(dataConsumingOffset);
+
+ int readSize =
+ readBuffer(
+ fileInfo.getFilePath(),
+ dataFileChannel,
+ headerBuffer,
+ buffer,
+ headerBuffer.capacity());
+ currentPartitionRemainingBytes -= readSize;
+
+ logger.debug(
+ "readBuffer data: {}, {}, {}, {}, {}, {}",
+ streamId,
+ currentPartitionRemainingBytes,
+ readSize,
+ dataConsumingOffset,
+ fileInfo.getFilePath(),
+ System.identityHashCode(buffer));
+
+ // if this check fails, the partition file must be corrupted
+ if (currentPartitionRemainingBytes < 0) {
+ throw new RuntimeException("File is corrupted");
+ } else if (currentPartitionRemainingBytes == 0) {
+ logger.debug(
+ "readBuffer end, {}, {}, {}, {}",
+ streamId,
+ dataFileChannel.size(),
+ dataConsumingOffset,
+ currentPartitionRemainingBytes);
+ int prevDataRegion = currentDataRegion;
+ updateConsumingOffset();
+ return prevDataRegion == currentDataRegion &&
currentPartitionRemainingBytes > 0;
+ }
+
+ dataConsumingOffset = dataFileChannel.position();
+
+ logger.debug(
+ "readBuffer run: {}, {}, {}, {}",
+ streamId,
+ dataFileChannel.size(),
+ dataConsumingOffset,
+ currentPartitionRemainingBytes);
+ return true;
+ } catch (Throwable throwable) {
+ logger.debug("Failed to read partition file.", throwable);
+ isReleased = true;
+ throw throwable;
+ }
+ }
+
+ public boolean hasRemaining() {
+ return currentPartitionRemainingBytes > 0;
+ }
+
+ private void notifyBacklog(int backlog) {
+ logger.debug("stream manager stream id {} backlog:{}", streamId, backlog);
+ associatedChannel.writeAndFlush(new BacklogAnnouncement(streamId,
backlog));
+ }
+
+ private void notifyError(Throwable throwable) {
+ logger.error("read error stream id {} message:{}", streamId,
throwable.getMessage(), throwable);
+ // TODO notify client the exception
+ }
+
+ public long getPriority() {
+ return dataConsumingOffset;
+ }
+
+ @Override
+ public int compareTo(DataPartitionReader that) {
+ return Long.compare(getPriority(), that.getPriority());
+ }
+
+ public FileInfo getFileInfo() {
+ return fileInfo;
+ }
+
+ public void closeReader() {
+ synchronized (lock) {
+ isClosed = true;
+ }
+
+ logger.debug("Closed read for stream {}", this.streamId);
+ }
+
+ private void recycle() {
+ synchronized (lock) {
+ if (!isReleased) {
+ release();
+ recycleStream.run();
+ }
+ }
+ }
+
+ public void recycleOnError(Throwable throwable) {
+ synchronized (lock) {
+ if (!isError) {
+ isError = true;
+ errorCause = throwable;
+ notifyError(throwable);
+ recycle();
+ }
+ }
+ }
+
+ public void release() {
+ // we can safely release if reader reaches error or (read/send finished)
+ synchronized (lock) {
+ if (!isReleased) {
+ logger.debug("release reader for stream {}", this.streamId);
+ isReleased = true;
+ if (!buffersRead.isEmpty()) {
+ buffersRead.forEach(WrappedDataBuffer::release);
+ }
+ }
+ }
+ }
+
+ public boolean isFinished() {
+ synchronized (lock) {
+ // ensure every buffer are return to bufferQueue or release in
buffersRead
+ return numInFlightRequests.get() == 0 && isReleased;
+ }
+ }
+
+ @Override
+ public String toString() {
+ final StringBuilder sb = new StringBuilder("DataPartitionReader{");
+ sb.append("startPartitionIndex=").append(startPartitionIndex);
+ sb.append(", endPartitionIndex=").append(endPartitionIndex);
+ sb.append(", streamId=").append(streamId);
+ sb.append('}');
+ return sb.toString();
+ }
+}
diff --git
a/common/src/main/java/org/apache/celeborn/common/network/server/memory/BufferRecycler.java
b/common/src/main/java/org/apache/celeborn/common/network/server/memory/BufferRecycler.java
new file mode 100644
index 000000000..069bc81e0
--- /dev/null
+++
b/common/src/main/java/org/apache/celeborn/common/network/server/memory/BufferRecycler.java
@@ -0,0 +1,44 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.common.network.server.memory;
+
+import java.util.function.Consumer;
+
+import io.netty.buffer.ByteBuf;
+
+public class BufferRecycler implements Recycler {
+
+ private MemoryManager memoryManager;
+
+ private Consumer<ByteBuf> recycleConsumer;
+
+ public BufferRecycler(MemoryManager memoryManager, Consumer<ByteBuf>
recycleConsumer) {
+ this.memoryManager = memoryManager;
+ this.recycleConsumer = recycleConsumer;
+ }
+
+ @Override
+ public void recycle(ByteBuf byteBuf) {
+ recycleConsumer.accept(byteBuf);
+ }
+
+ @Override
+ public void release(ByteBuf byteBuf) {
+ memoryManager.recycleReadBuffer(byteBuf);
+ }
+}
diff --git
a/common/src/main/java/org/apache/celeborn/common/network/server/memory/ReadBufferDispatcher.java
b/common/src/main/java/org/apache/celeborn/common/network/server/memory/ReadBufferDispatcher.java
index ca76a25f8..ecaddb72d 100644
---
a/common/src/main/java/org/apache/celeborn/common/network/server/memory/ReadBufferDispatcher.java
+++
b/common/src/main/java/org/apache/celeborn/common/network/server/memory/ReadBufferDispatcher.java
@@ -30,7 +30,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.network.util.NettyUtils;
public class ReadBufferDispatcher extends Thread {
- private Logger logger = LoggerFactory.getLogger(ReadBufferDispatcher.class);
+ private final Logger logger =
LoggerFactory.getLogger(ReadBufferDispatcher.class);
private final LinkedBlockingQueue<ReadBufferRequest> requests = new
LinkedBlockingQueue<>();
private final MemoryManager memoryManager;
private final PooledByteBufAllocator readBufferAllocator;
@@ -48,7 +48,11 @@ public class ReadBufferDispatcher extends Thread {
public void recycle(ByteBuf buf) {
int bufferSize = buf.capacity();
- buf.release();
+ if (buf.refCnt() == 0) {
+ logger.warn("recycle encounter: {}", buf.refCnt());
+ } else {
+ buf.release(buf.refCnt());
+ }
memoryManager.changeReadBufferCounter(-1 * bufferSize);
}
@@ -61,6 +65,7 @@ public class ReadBufferDispatcher extends Thread {
} catch (InterruptedException e) {
logger.info("Buffer dispatcher is closing");
}
+
if (request != null) {
List<ByteBuf> buffers = new ArrayList<>();
int bufferSize = request.getBufferSize();
@@ -68,7 +73,6 @@ public class ReadBufferDispatcher extends Thread {
if (memoryManager.readBufferAvailable(bufferSize)) {
memoryManager.incrementDiskBuffer(bufferSize);
ByteBuf buf = readBufferAllocator.buffer(bufferSize, bufferSize);
- buf.retain();
buffers.add(buf);
} else {
try {
@@ -86,7 +90,6 @@ public class ReadBufferDispatcher extends Thread {
&& buffers.size() < request.getMax()) {
memoryManager.changeReadBufferCounter(bufferSize);
ByteBuf buf = readBufferAllocator.buffer(bufferSize, bufferSize);
- buf.retain();
buffers.add(buf);
}
request.getBufferListener().notifyBuffers(buffers, null);
diff --git
a/common/src/main/java/org/apache/celeborn/common/network/server/memory/Recycler.java
b/common/src/main/java/org/apache/celeborn/common/network/server/memory/Recycler.java
new file mode 100644
index 000000000..2c8493d7f
--- /dev/null
+++
b/common/src/main/java/org/apache/celeborn/common/network/server/memory/Recycler.java
@@ -0,0 +1,30 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.common.network.server.memory;
+
+import io.netty.buffer.ByteBuf;
+
+/** manager bytebuf lifecycle */
+public interface Recycler {
+
+ /** recycle bytebuf to buffer pool */
+ void recycle(ByteBuf byteBuf);
+
+ /** release the bytebuf */
+ void release(ByteBuf byteBuf);
+}
diff --git
a/common/src/main/java/org/apache/celeborn/common/network/server/memory/WrappedDataBuffer.java
b/common/src/main/java/org/apache/celeborn/common/network/server/memory/WrappedDataBuffer.java
new file mode 100644
index 000000000..fdbd9d1be
--- /dev/null
+++
b/common/src/main/java/org/apache/celeborn/common/network/server/memory/WrappedDataBuffer.java
@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.celeborn.common.network.server.memory;
+
+import io.netty.buffer.ByteBuf;
+
+public class WrappedDataBuffer {
+
+ public final ByteBuf byteBuf;
+
+ public final Recycler bufferRecycler;
+
+ public WrappedDataBuffer(ByteBuf byteBuf, Recycler bufferRecycler) {
+ this.byteBuf = byteBuf;
+ this.bufferRecycler = bufferRecycler;
+ }
+
+ public void recycle() {
+ bufferRecycler.recycle(byteBuf);
+ }
+
+ public void release() {
+ bufferRecycler.release(byteBuf);
+ }
+}
