[FLINK-25796][network] Avoid record copy for result partition of sort-shuffle if there are enough buffers for better performance

Author: kevin.cyj

flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/BufferBuilder.java

@@ -44,13 +44,8 @@ public class BufferBuilder implements AutoCloseable {
     private boolean bufferConsumerCreated = false;
 
     public BufferBuilder(MemorySegment memorySegment, BufferRecycler recycler) {
-        this(memorySegment, Buffer.DataType.DATA_BUFFER, recycler);
-    }
-
-    public BufferBuilder(
-            MemorySegment memorySegment, Buffer.DataType dataType, BufferRecycler recycler) {
         this.memorySegment = checkNotNull(memorySegment);
-        this.buffer = new NetworkBuffer(memorySegment, recycler, dataType);
+        this.buffer = new NetworkBuffer(memorySegment, recycler);
         this.maxCapacity = buffer.getMaxCapacity();
     }
 

flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/SparsePartitionSortedBuffer.java renamed to flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/HashBasedPartitionSortedBuffer.java

@@ -24,20 +24,29 @@
 import org.apache.flink.runtime.io.network.buffer.BufferBuilder;
 import org.apache.flink.runtime.io.network.buffer.BufferConsumer;
 import org.apache.flink.runtime.io.network.buffer.BufferPool;
+import org.apache.flink.runtime.io.network.buffer.FreeingBufferRecycler;
+import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
 
 import javax.annotation.Nullable;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.ArrayDeque;
-import java.util.Queue;
 
 import static org.apache.flink.util.Preconditions.checkArgument;
 import static org.apache.flink.util.Preconditions.checkNotNull;
 import static org.apache.flink.util.Preconditions.checkState;
 
-/** Placeholder. */
-public class SparsePartitionSortedBuffer implements SortBuffer {
+/**
+ * * A {@link SortBuffer} implementation which sorts all appended records only by subpartition
+ * index. Records of the same subpartition keep the appended order.
+ *
+ * <p>Different from the {@link SortBasedPartitionSortedBuffer}, in this {@link SortBuffer}
+ * implementation, memory segment boundary serves as the nature data boundary of different
+ * subpartitions, which means that one memory segment can never contain data from different
+ * subpartitions.
+ */
+public class HashBasedPartitionSortedBuffer implements SortBuffer {
 
     /** A buffer pool to request memory segments from. */
     private final BufferPool bufferPool;
@@ -46,40 +55,51 @@ public class SparsePartitionSortedBuffer implements SortBuffer {
     private final int numGuaranteedBuffers;
 
     /** Buffers containing data for all subpartitions. */
-    private final ArrayDeque<BufferBuilder>[] buffers;
+    private final ArrayDeque<BufferConsumer>[] buffers;
 
     // ---------------------------------------------------------------------------------------------
     // Statistics and states
     // ---------------------------------------------------------------------------------------------
 
-    private int numTotalBuffers;
-
     /** Total number of bytes already appended to this sort buffer. */
     private long numTotalBytes;
 
     /** Total number of records already appended to this sort buffer. */
     private long numTotalRecords;
 
+    /** Whether this sort buffer is full and ready to read data from. */
+    private boolean isFull;
+
     /** Whether this sort buffer is finished. One can only read a finished sort buffer. */
     private boolean isFinished;
 
     /** Whether this sort buffer is released. A released sort buffer can not be used. */
     private boolean isReleased;
 
+    // ---------------------------------------------------------------------------------------------
+    // For writing
+    // ---------------------------------------------------------------------------------------------
+
+    /** Partial buffers to be appended data for each channel. */
+    private final BufferBuilder[] builders;
+
+    /** Total number of network buffers already occupied currently by this sort buffer. */
+    private int numBuffersOccupied;
+
     // ---------------------------------------------------------------------------------------------
     // For reading
     // ---------------------------------------------------------------------------------------------
 
+    /** Used to index the current available channel to read data from. */
+    private int readOrderIndex;
+
     /** Data of different subpartitions in this sort buffer will be read in this order. */
     private final int[] subpartitionReadOrder;
 
     /** Total number of bytes already read from this sort buffer. */
     private long numTotalBytesRead;
 
-    /** Used to index the current available channel to read data from. */
-    private int readOrderIndex;
-
-    public SparsePartitionSortedBuffer(
+    public HashBasedPartitionSortedBuffer(
             BufferPool bufferPool,
             int numSubpartitions,
             int numGuaranteedBuffers,
@@ -89,6 +109,7 @@ public SparsePartitionSortedBuffer(
         this.bufferPool = checkNotNull(bufferPool);
         this.numGuaranteedBuffers = numGuaranteedBuffers;
 
+        this.builders = new BufferBuilder[numSubpartitions];
         this.buffers = new ArrayDeque[numSubpartitions];
         for (int channel = 0; channel < numSubpartitions; ++channel) {
             this.buffers[channel] = new ArrayDeque<>();
@@ -109,133 +130,114 @@ public SparsePartitionSortedBuffer(
     public boolean append(ByteBuffer source, int targetChannel, Buffer.DataType dataType)
             throws IOException {
         checkArgument(source.hasRemaining(), "Cannot append empty data.");
+        checkState(!isFull, "Sort buffer is already full.");
         checkState(!isFinished, "Sort buffer is already finished.");
         checkState(!isReleased, "Sort buffer is already released.");
 
         int totalBytes = source.remaining();
         if (dataType.isBuffer()) {
-            // return false directly if it can not allocate enough buffers for the given record
-            Queue<BufferBuilder> freeSegments = allocateBuffersForRecord(totalBytes, targetChannel);
-            if (freeSegments == null) {
-                return false;
-            }
-            writeRecord(source, targetChannel, freeSegments);
+            writeRecord(source, targetChannel);
         } else {
             writeEvent(source, targetChannel, dataType);
         }
 
-        ++numTotalRecords;
-        numTotalBytes += totalBytes;
-        return true;
+        isFull = source.hasRemaining();
+        if (!isFull) {
+            ++numTotalRecords;
+        }
+        numTotalBytes += totalBytes - source.remaining();
+        return isFull;
     }
 
     private void writeEvent(ByteBuffer source, int targetChannel, Buffer.DataType dataType) {
-        ArrayDeque<BufferBuilder> channelBuffers = buffers[targetChannel];
-        if (!channelBuffers.isEmpty()) {
-            channelBuffers.peekLast().finish();
+        BufferBuilder builder = builders[targetChannel];
+        if (builder != null) {
+            builder.finish();
+            buffers[targetChannel].add(builder.createBufferConsumerFromBeginning());
+            builder.close();
+            builders[targetChannel] = null;
         }
 
-        MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(source.remaining());
-        BufferBuilder builder = new BufferBuilder(segment, dataType, ignored -> {});
-        builder.append(source);
-        channelBuffers.add(builder);
+        MemorySegment segment =
+                MemorySegmentFactory.allocateUnpooledOffHeapMemory(source.remaining());
+        segment.put(0, source, segment.size());
+        BufferConsumer consumer =
+                new BufferConsumer(
+                        new NetworkBuffer(segment, FreeingBufferRecycler.INSTANCE, dataType),
+                        segment.size());
+        buffers[targetChannel].add(consumer);
     }
 
-    private void writeRecord(
-            ByteBuffer source, int targetChannel, Queue<BufferBuilder> freeSegments) {
-        int remainingFreeSpace = 0;
-        ArrayDeque<BufferBuilder> channelBuffers = buffers[targetChannel];
-
-        BufferBuilder prevBuffer = null;
-        if (!channelBuffers.isEmpty()) {
-            prevBuffer = channelBuffers.peekLast();
-            remainingFreeSpace = prevBuffer.getWritableBytes();
-        }
-
-        if (remainingFreeSpace > 0) {
-            prevBuffer.append(source);
-        }
-
-        while (source.hasRemaining()) {
-            BufferBuilder targetBuffer = freeSegments.poll();
-            targetBuffer.append(source);
-            if (prevBuffer != null) {
-                prevBuffer.finish();
+    private void writeRecord(ByteBuffer source, int targetChannel) throws IOException {
+        do {
+            BufferBuilder builder = builders[targetChannel];
+            if (builder == null) {
+                builder = requestBufferFromPool();
+                if (builder == null) {
+                    break;
+                }
+                ++numBuffersOccupied;
+                builders[targetChannel] = builder;
             }
-            channelBuffers.add(targetBuffer);
-            prevBuffer = targetBuffer;
-        }
-    }
-
-    private ArrayDeque<BufferBuilder> allocateBuffersForRecord(
-            int numRecordBytes, int targetChannel) throws IOException {
-        int remainingFreeSpace = 0;
-        ArrayDeque<BufferBuilder> channelBuffers = buffers[targetChannel];
-        if (!channelBuffers.isEmpty()) {
-            remainingFreeSpace = channelBuffers.peekLast().getWritableBytes();
-        }
 
-        ArrayDeque<BufferBuilder> freeBuffers = new ArrayDeque<>();
-        while (remainingFreeSpace < numRecordBytes) {
-            BufferBuilder buffer = requestBufferFromPool();
-            if (buffer == null) {
-                // return null if we can not allocate enough buffers for the appended record
-                numTotalBuffers -= freeBuffers.size();
-                freeBuffers.forEach(BufferBuilder::close);
-                freeBuffers.clear();
-                return null;
+            builder.append(source);
+            if (builder.isFull()) {
+                builder.finish();
+                buffers[targetChannel].add(builder.createBufferConsumerFromBeginning());
+                builder.close();
+                builders[targetChannel] = null;
             }
-
-            ++numTotalBuffers;
-            freeBuffers.add(buffer);
-            remainingFreeSpace += buffer.getMaxCapacity();
-        }
-        return freeBuffers;
+        } while (source.hasRemaining());
     }
 
     private BufferBuilder requestBufferFromPool() throws IOException {
         try {
             // blocking request buffers if there is still guaranteed memory
-            if (numTotalBuffers < numGuaranteedBuffers) {
+            if (numBuffersOccupied < numGuaranteedBuffers) {
                 return bufferPool.requestBufferBuilderBlocking();
             }
         } catch (InterruptedException e) {
-            throw new IOException("Interrupted while requesting buffer.");
+            throw new IOException("Interrupted while requesting buffer.", e);
         }
 
         return bufferPool.requestBufferBuilder();
     }
 
     @Override
     public BufferWithChannel copyIntoSegment(MemorySegment target) {
-        checkState(hasRemaining(), "No data remaining.");
-        checkState(isFinished, "Should finish the sort buffer first before coping any data.");
+        checkState(isFull, "Sort buffer is not ready to be read.");
         checkState(!isReleased, "Sort buffer is already released.");
 
         BufferWithChannel buffer = null;
+        if (!hasRemaining() || readOrderIndex >= subpartitionReadOrder.length) {
+            return null;
+        }
+
         int targetChannel = subpartitionReadOrder[readOrderIndex];
         while (buffer == null) {
-            try (BufferBuilder builder = buffers[targetChannel].poll()) {
-                if (builder != null) {
-                    try (BufferConsumer consumer = builder.createBufferConsumerFromBeginning()) {
-                        buffer = new BufferWithChannel(consumer.build(), targetChannel);
-                        numTotalBytesRead += buffer.getBuffer().readableBytes();
-                    }
-                } else {
-                    targetChannel = subpartitionReadOrder[++readOrderIndex];
+            BufferConsumer consumer = buffers[targetChannel].poll();
+            if (consumer != null) {
+                buffer = new BufferWithChannel(consumer.build(), targetChannel);
+                numBuffersOccupied -= buffer.getBuffer().isBuffer() ? 1 : 0;
+                numTotalBytesRead += buffer.getBuffer().readableBytes();
+                consumer.close();
+            } else {
+                if (++readOrderIndex >= subpartitionReadOrder.length) {
+                    break;
                 }
+                targetChannel = subpartitionReadOrder[readOrderIndex];
             }
         }
         return buffer;
     }
 
     @Override
-    public long numRecords() {
+    public long numTotalRecords() {
         return numTotalRecords;
     }
 
     @Override
-    public long numBytes() {
+    public long numTotalBytes() {
         return numTotalBytes;
     }
 
@@ -246,12 +248,18 @@ public boolean hasRemaining() {
 
     @Override
     public void finish() {
+        checkState(!isFull, "SortBuffer must not be full.");
         checkState(!isFinished, "SortBuffer is already finished.");
 
+        isFull = true;
         isFinished = true;
-        for (ArrayDeque<BufferBuilder> channelBuffers : buffers) {
-            if (!channelBuffers.isEmpty()) {
-                channelBuffers.peekLast().finish();
+        for (int channel = 0; channel < builders.length; ++channel) {
+            BufferBuilder builder = builders[channel];
+            if (builder != null) {
+                builder.finish();
+                buffers[channel].add(builder.createBufferConsumerFromBeginning());
+                builder.close();
+                builders[channel] = null;
             }
         }
     }
@@ -263,13 +271,26 @@ public boolean isFinished() {
 
     @Override
     public void release() {
-        // the sort buffer can be released by other threads
         if (isReleased) {
             return;
         }
         isReleased = true;
 
-        clear();
+        for (int channel = 0; channel < builders.length; ++channel) {
+            BufferBuilder builder = builders[channel];
+            if (builder != null) {
+                builder.close();
+                builders[channel] = null;
+            }
+        }
+
+        for (ArrayDeque<BufferConsumer> buffer : buffers) {
+            BufferConsumer consumer = buffer.poll();
+            while (consumer != null) {
+                consumer.close();
+                consumer = buffer.poll();
+            }
+        }
     }
 
     @Override
@@ -278,18 +299,16 @@ public boolean isReleased() {
     }
 
     @Override
-    public void clear() {
-        for (ArrayDeque<BufferBuilder> channelBuffers : buffers) {
-            while (!channelBuffers.isEmpty()) {
-                channelBuffers.poll().close();
-            }
-        }
+    public void reset() {
+        checkState(!isFinished, "Sort buffer has been finished.");
+        checkState(!isReleased, "Sort buffer has been released.");
 
-        numTotalBuffers = 0;
-        numTotalBytes = 0;
-        numTotalRecords = 0;
-        isFinished = false;
-        numTotalBytesRead = 0;
+        isFull = false;
         readOrderIndex = 0;
     }
+
+    @Override
+    public BufferPool getBufferPool() {
+        return bufferPool;
+    }
 }