This is an automated email from the ASF dual-hosted git repository.
guoweijie pushed a commit to branch release-1.18
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/release-1.18 by this push:
new 1b3849c196c [FLINK-35169][runtime] Recycle buffers to freeSegments
before releasing data buffer for sort accumulator
1b3849c196c is described below
commit 1b3849c196c25784dc662398385e0b30f2c23a03
Author: Yuxin Tan <[email protected]>
AuthorDate: Fri Apr 19 16:40:55 2024 +0800
[FLINK-35169][runtime] Recycle buffers to freeSegments before releasing
data buffer for sort accumulator
---
.../tiered/storage/SortBufferAccumulator.java | 92 +++++++++++++++-------
.../tiered/storage/SortBufferAccumulatorTest.java | 44 +++++++++++
2 files changed, 108 insertions(+), 28 deletions(-)
diff --git
a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java
b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java
index fc0d1a56f5e..4c7e458f76a 100644
---
a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java
+++
b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulator.java
@@ -29,6 +29,7 @@ import
org.apache.flink.runtime.io.network.partition.DataBuffer;
import
org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
import javax.annotation.Nullable;
+import javax.annotation.concurrent.GuardedBy;
import java.io.IOException;
import java.nio.ByteBuffer;
@@ -67,6 +68,7 @@ public class SortBufferAccumulator implements
BufferAccumulator {
private final int bufferSizeBytes;
/** The empty buffers without storing data. */
+ @GuardedBy("lock")
private final LinkedList<MemorySegment> freeSegments = new LinkedList<>();
/** The memory manager of the tiered storage. */
@@ -78,7 +80,9 @@ public class SortBufferAccumulator implements
BufferAccumulator {
* data integrity. Note that this can be null before using it to store
records, and this {@link
* DataBuffer} will be released once flushed.
*/
- @Nullable private DataBuffer currentDataBuffer;
+ @GuardedBy("lock")
+ @Nullable
+ private DataBuffer currentDataBuffer;
/**
* The buffer recycler. Note that this can be null before requesting
buffers from the memory
@@ -97,6 +101,11 @@ public class SortBufferAccumulator implements
BufferAccumulator {
/** Whether the current {@link DataBuffer} is a broadcast sort buffer. */
private boolean isBroadcastDataBuffer;
+ @GuardedBy("lock")
+ private boolean isDataBufferReleased;
+
+ private final Object lock = new Object();
+
public SortBufferAccumulator(
int numSubpartitions,
int numBuffers,
@@ -121,38 +130,44 @@ public class SortBufferAccumulator implements
BufferAccumulator {
boolean isBroadcast)
throws IOException {
int targetSubpartition = subpartitionId.getSubpartitionId();
- switchCurrentDataBufferIfNeeded(isBroadcast);
- if (!checkNotNull(currentDataBuffer).append(record,
targetSubpartition, dataType)) {
- return;
- }
+ synchronized (lock) {
+ switchCurrentDataBufferIfNeeded(isBroadcast);
+ if (!checkNotNull(currentDataBuffer).append(record,
targetSubpartition, dataType)) {
+ return;
+ }
- // The sort buffer is empty, but we failed to write the record into
it, which indicates the
- // record is larger than the sort buffer can hold. So the record is
written into multiple
- // buffers directly.
- if (!currentDataBuffer.hasRemaining()) {
- currentDataBuffer.release();
- writeLargeRecord(record, targetSubpartition, dataType);
- return;
+ // The sort buffer is empty, but we failed to write the record
into it, which indicates
+ // the record is larger than the sort buffer can hold. So the
record is written into
+ // multiple buffers directly.
+ if (!currentDataBuffer.hasRemaining()) {
+ isDataBufferReleased = true;
+ currentDataBuffer.release();
+ writeLargeRecord(record, targetSubpartition, dataType);
+ return;
+ }
+ flushDataBuffer();
}
- flushDataBuffer();
checkState(record.hasRemaining(), "Empty record.");
receive(record, subpartitionId, dataType, isBroadcast);
}
@Override
public void close() {
- flushCurrentDataBuffer();
- releaseFreeBuffers();
- if (currentDataBuffer != null) {
- currentDataBuffer.release();
+ synchronized (lock) {
+ flushCurrentDataBuffer();
+ isDataBufferReleased = true;
+ releaseFreeBuffers();
+ if (currentDataBuffer != null) {
+ currentDataBuffer.release();
+ }
}
}
// ------------------------------------------------------------------------
// Internal Methods
// ------------------------------------------------------------------------
-
+ @GuardedBy("lock")
private void switchCurrentDataBufferIfNeeded(boolean isBroadcast) {
if (isBroadcast == isBroadcastDataBuffer
&& currentDataBuffer != null
@@ -162,9 +177,12 @@ public class SortBufferAccumulator implements
BufferAccumulator {
}
isBroadcastDataBuffer = isBroadcast;
flushCurrentDataBuffer();
+ isDataBufferReleased = true;
currentDataBuffer = createNewDataBuffer();
+ isDataBufferReleased = false;
}
+ @GuardedBy("lock")
private DataBuffer createNewDataBuffer() {
requestBuffers();
@@ -178,6 +196,7 @@ public class SortBufferAccumulator implements
BufferAccumulator {
numBuffersForSort);
}
+ @GuardedBy("lock")
private void requestBuffers() {
while (freeSegments.size() < numBuffers) {
Buffer buffer = requestBuffer();
@@ -188,6 +207,7 @@ public class SortBufferAccumulator implements
BufferAccumulator {
}
}
+ @GuardedBy("lock")
private void flushDataBuffer() {
if (currentDataBuffer == null
|| currentDataBuffer.isReleased()
@@ -205,14 +225,17 @@ public class SortBufferAccumulator implements
BufferAccumulator {
flushBuffer(bufferWithChannel);
} while (true);
+ isDataBufferReleased = true;
releaseFreeBuffers();
currentDataBuffer.release();
}
private void flushCurrentDataBuffer() {
- if (currentDataBuffer != null) {
- flushDataBuffer();
- currentDataBuffer = null;
+ synchronized (lock) {
+ if (currentDataBuffer != null) {
+ flushDataBuffer();
+ currentDataBuffer = null;
+ }
}
}
@@ -235,11 +258,13 @@ public class SortBufferAccumulator implements
BufferAccumulator {
}
private MemorySegment getFreeSegment() {
- MemorySegment freeSegment = freeSegments.poll();
- if (freeSegment == null) {
- freeSegment = requestBuffer().getMemorySegment();
+ synchronized (lock) {
+ MemorySegment freeSegment = freeSegments.poll();
+ if (freeSegment == null) {
+ freeSegment = requestBuffer().getMemorySegment();
+ }
+ return freeSegment;
}
- return freeSegment;
}
private void flushBuffer(BufferWithChannel bufferWithChannel) {
@@ -259,11 +284,22 @@ public class SortBufferAccumulator implements
BufferAccumulator {
}
private void releaseFreeBuffers() {
- freeSegments.forEach(this::recycleBuffer);
- freeSegments.clear();
+ synchronized (lock) {
+ isDataBufferReleased = true;
+ freeSegments.forEach(this::recycleBuffer);
+ freeSegments.clear();
+ }
}
private void recycleBuffer(MemorySegment memorySegment) {
- checkNotNull(bufferRecycler).recycle(memorySegment);
+ synchronized (lock) {
+ if (!isDataBufferReleased
+ && currentDataBuffer != null
+ && !currentDataBuffer.isReleased()) {
+ freeSegments.add(memorySegment);
+ } else {
+ checkNotNull(bufferRecycler).recycle(memorySegment);
+ }
+ }
}
}
diff --git
a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java
b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java
index a00fd705b90..228b05119c3 100644
---
a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java
+++
b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/tiered/storage/SortBufferAccumulatorTest.java
@@ -18,9 +18,12 @@
package org.apache.flink.runtime.io.network.partition.hybrid.tiered.storage;
+import org.apache.flink.core.memory.MemorySegment;
import org.apache.flink.runtime.io.network.buffer.Buffer;
+import org.apache.flink.runtime.io.network.buffer.BufferBuilder;
import org.apache.flink.runtime.io.network.buffer.BufferPool;
import org.apache.flink.runtime.io.network.buffer.NetworkBufferPool;
+import
org.apache.flink.runtime.io.network.partition.hybrid.tiered.TestingTieredStorageMemoryManager;
import
org.apache.flink.runtime.io.network.partition.hybrid.tiered.common.TieredStorageSubpartitionId;
import org.junit.jupiter.api.AfterEach;
@@ -174,6 +177,47 @@ class SortBufferAccumulatorTest {
assertThat(tieredStorageMemoryManager.numOwnerRequestedBuffer(bufferAccumulator)).isZero();
}
+ @Test
+ void testReuseRecycledBuffersWhenFlushDataBuffer() throws IOException {
+ int numBuffers = 10;
+ int numSubpartitions = 10;
+
+ AtomicInteger numRequestedBuffers = new AtomicInteger(0);
+ TestingTieredStorageMemoryManager memoryManager =
+ new TestingTieredStorageMemoryManager.Builder()
+ .setRequestBufferBlockingFunction(
+ owner -> {
+ numRequestedBuffers.incrementAndGet();
+ MemorySegment memorySegment =
+
globalPool.requestPooledMemorySegment();
+ // When flushing data from the data
buffer, the buffers should
+ // not be requested from here anymore
because the freeSegments
+ // can be reused before released.
+ assertThat(numRequestedBuffers.get() <=
numBuffers).isTrue();
+ return new BufferBuilder(
+ memorySegment,
+ segment ->
globalPool.requestPooledMemorySegment());
+ })
+ .build();
+
+ SortBufferAccumulator bufferAccumulator =
+ new SortBufferAccumulator(
+ numSubpartitions, numBuffers, BUFFER_SIZE_BYTES,
memoryManager);
+ bufferAccumulator.setup(
+ ((subpartition, buffers) ->
buffers.forEach(Buffer::recycleBuffer)));
+
+ for (int i = 0; i < numSubpartitions; i++) {
+ bufferAccumulator.receive(
+ generateRandomData(10, new Random()),
+ new TieredStorageSubpartitionId(i % numSubpartitions),
+ Buffer.DataType.DATA_BUFFER,
+ false);
+ }
+
+ // Trigger flushing buffers from data buffer.
+ bufferAccumulator.close();
+ }
+
private TieredStorageMemoryManagerImpl createStorageMemoryManager(int
numBuffersInBufferPool)
throws IOException {
BufferPool bufferPool =
