You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by xt...@apache.org on 2022/09/01 03:35:37 UTC
[flink] branch master updated: [FLINK-29053] Hybrid shuffle has concurrent modification of buffer when compression is enabled
This is an automated email from the ASF dual-hosted git repository.
xtsong pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push:
new 8b8245ba46b [FLINK-29053] Hybrid shuffle has concurrent modification of buffer when compression is enabled
8b8245ba46b is described below
commit 8b8245ba46b25c2617d91cff3d3a44b99879d9f2
Author: Weijie Guo <re...@163.com>
AuthorDate: Fri Aug 19 18:51:19 2022 +0800
[FLINK-29053] Hybrid shuffle has concurrent modification of buffer when compression is enabled
This closes #20647
---
.../runtime/io/network/buffer/NetworkBuffer.java | 6 +-
.../buffer/ReadOnlySlicedNetworkBuffer.java | 10 +-
.../partition/hybrid/HsMemoryDataManager.java | 8 +-
.../partition/hybrid/HsMemoryDataSpiller.java | 35 +-----
.../hybrid/HsSubpartitionMemoryDataManager.java | 24 +++-
.../partition/hybrid/HsMemoryDataSpillerTest.java | 46 ++++----
.../HsSubpartitionMemoryDataManagerTest.java | 121 +++++++++++++++------
7 files changed, 151 insertions(+), 99 deletions(-)
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/NetworkBuffer.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/NetworkBuffer.java
index ef094e50d06..4c48e12fe8a 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/NetworkBuffer.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/NetworkBuffer.java
@@ -173,8 +173,10 @@ public class NetworkBuffer extends AbstractReferenceCountedByteBuf implements Bu
@Override
public ReadOnlySlicedNetworkBuffer readOnlySlice(int index, int length) {
- checkState(!isCompressed, "Unable to slice a compressed buffer.");
- return new ReadOnlySlicedNetworkBuffer(this, index, length);
+ checkState(
+ !isCompressed || index + length != writerIndex(),
+ "Unable to partially slice a compressed buffer.");
+ return new ReadOnlySlicedNetworkBuffer(this, index, length, isCompressed);
}
@Override
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/ReadOnlySlicedNetworkBuffer.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/ReadOnlySlicedNetworkBuffer.java
index 4fccec09bdc..b40c4bea0a3 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/ReadOnlySlicedNetworkBuffer.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/ReadOnlySlicedNetworkBuffer.java
@@ -56,11 +56,13 @@ public final class ReadOnlySlicedNetworkBuffer extends ReadOnlyByteBuf implement
* @param buffer the buffer to derive from
* @param index the index to start from
* @param length the length of the slice
+ * @param isCompressed is the buffer compressed
*/
- ReadOnlySlicedNetworkBuffer(NetworkBuffer buffer, int index, int length) {
+ ReadOnlySlicedNetworkBuffer(NetworkBuffer buffer, int index, int length, boolean isCompressed) {
super(new SlicedByteBuf(buffer, index, length));
this.memorySegmentOffset = buffer.getMemorySegmentOffset() + index;
this.dataType = buffer.getDataType();
+ this.isCompressed = isCompressed;
}
/**
@@ -141,9 +143,11 @@ public final class ReadOnlySlicedNetworkBuffer extends ReadOnlyByteBuf implement
@Override
public ReadOnlySlicedNetworkBuffer readOnlySlice(int index, int length) {
- checkState(!isCompressed, "Unable to slice a compressed buffer.");
+ checkState(
+ !isCompressed || index + length == writerIndex(),
+ "Unable to partially slice a compressed buffer.");
return new ReadOnlySlicedNetworkBuffer(
- super.unwrap(), index, length, memorySegmentOffset, false);
+ super.unwrap(), index, length, memorySegmentOffset, isCompressed);
}
@Override
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataManager.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataManager.java
index 2fc63646c03..867ed539d72 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataManager.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataManager.java
@@ -82,7 +82,7 @@ public class HsMemoryDataManager implements HsSpillingInfoProvider, HsMemoryData
throws IOException {
this.numSubpartitions = numSubpartitions;
this.bufferPool = bufferPool;
- this.spiller = new HsMemoryDataSpiller(dataFilePath, bufferCompressor);
+ this.spiller = new HsMemoryDataSpiller(dataFilePath);
this.spillStrategy = spillStrategy;
this.fileDataIndex = fileDataIndex;
this.subpartitionMemoryDataManagers = new HsSubpartitionMemoryDataManager[numSubpartitions];
@@ -93,7 +93,11 @@ public class HsMemoryDataManager implements HsSpillingInfoProvider, HsMemoryData
for (int subpartitionId = 0; subpartitionId < numSubpartitions; ++subpartitionId) {
subpartitionMemoryDataManagers[subpartitionId] =
new HsSubpartitionMemoryDataManager(
- subpartitionId, bufferSize, readWriteLock.readLock(), this);
+ subpartitionId,
+ bufferSize,
+ readWriteLock.readLock(),
+ bufferCompressor,
+ this);
}
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataSpiller.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataSpiller.java
index b6a8a9b67c0..e2f4cb06aba 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataSpiller.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataSpiller.java
@@ -19,7 +19,6 @@
package org.apache.flink.runtime.io.network.partition.hybrid;
import org.apache.flink.runtime.io.network.buffer.Buffer;
-import org.apache.flink.runtime.io.network.buffer.BufferCompressor;
import org.apache.flink.runtime.io.network.partition.BufferReaderWriterUtil;
import org.apache.flink.runtime.io.network.partition.hybrid.HsFileDataIndex.SpilledBuffer;
import org.apache.flink.util.ExceptionUtils;
@@ -27,8 +26,6 @@ import org.apache.flink.util.FatalExitExceptionHandler;
import org.apache.flink.shaded.guava30.com.google.common.util.concurrent.ThreadFactoryBuilder;
-import javax.annotation.Nullable;
-
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.FileChannel;
@@ -41,9 +38,6 @@ import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
-import java.util.stream.Collectors;
-
-import static org.apache.flink.util.Preconditions.checkNotNull;
/**
* This component is responsible for asynchronously writing in-memory data to disk. Each spilling
@@ -66,17 +60,13 @@ public class HsMemoryDataSpiller implements AutoCloseable {
/** File channel to write data. */
private final FileChannel dataFileChannel;
- @Nullable private final BufferCompressor bufferCompressor;
-
/** Records the current writing location. */
private long totalBytesWritten;
- public HsMemoryDataSpiller(Path dataFilePath, @Nullable BufferCompressor bufferCompressor)
- throws IOException {
+ public HsMemoryDataSpiller(Path dataFilePath) throws IOException {
this.dataFileChannel =
FileChannel.open(
dataFilePath, StandardOpenOption.CREATE_NEW, StandardOpenOption.WRITE);
- this.bufferCompressor = bufferCompressor;
}
/**
@@ -98,10 +88,6 @@ public class HsMemoryDataSpiller implements AutoCloseable {
List<BufferWithIdentity> toWrite,
CompletableFuture<List<SpilledBuffer>> spilledFuture) {
try {
- toWrite =
- toWrite.stream()
- .map(this::compressBuffersIfPossible)
- .collect(Collectors.toList());
List<SpilledBuffer> spilledBuffers = new ArrayList<>();
long expectedBytes = createSpilledBuffersAndGetTotalBytes(toWrite, spilledBuffers);
// write all buffers to file
@@ -192,23 +178,4 @@ public class HsMemoryDataSpiller implements AutoCloseable {
ExceptionUtils.rethrow(e);
}
}
-
- private BufferWithIdentity compressBuffersIfPossible(BufferWithIdentity bufferWithIdentity) {
- Buffer buffer = bufferWithIdentity.getBuffer();
- if (!canBeCompressed(buffer)) {
- return bufferWithIdentity;
- }
-
- buffer = checkNotNull(bufferCompressor).compressToOriginalBuffer(buffer);
- return new BufferWithIdentity(
- buffer, bufferWithIdentity.getBufferIndex(), bufferWithIdentity.getChannelIndex());
- }
-
- /**
- * Whether the buffer can be compressed or not. Note that event is not compressed because it is
- * usually small and the size can become even larger after compression.
- */
- private boolean canBeCompressed(Buffer buffer) {
- return bufferCompressor != null && buffer.isBuffer() && buffer.readableBytes() > 0;
- }
}
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManager.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManager.java
index 87b6c10bba1..d8040da0f56 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManager.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManager.java
@@ -24,6 +24,7 @@ import org.apache.flink.core.memory.MemorySegmentFactory;
import org.apache.flink.runtime.io.network.buffer.Buffer;
import org.apache.flink.runtime.io.network.buffer.Buffer.DataType;
import org.apache.flink.runtime.io.network.buffer.BufferBuilder;
+import org.apache.flink.runtime.io.network.buffer.BufferCompressor;
import org.apache.flink.runtime.io.network.buffer.BufferConsumer;
import org.apache.flink.runtime.io.network.buffer.FreeingBufferRecycler;
import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
@@ -84,17 +85,21 @@ public class HsSubpartitionMemoryDataManager implements HsDataView {
/** DO NOT USE DIRECTLY. Use {@link #runWithLock} or {@link #callWithLock} instead. */
private final Object subpartitionLock = new Object();
+ @Nullable private final BufferCompressor bufferCompressor;
+
@Nullable private HsOutputMetrics outputMetrics;
HsSubpartitionMemoryDataManager(
int targetChannel,
int bufferSize,
Lock resultPartitionLock,
+ @Nullable BufferCompressor bufferCompressor,
HsMemoryDataManagerOperation memoryDataManagerOperation) {
this.targetChannel = targetChannel;
this.bufferSize = bufferSize;
this.resultPartitionLock = resultPartitionLock;
this.memoryDataManagerOperation = memoryDataManagerOperation;
+ this.bufferCompressor = bufferCompressor;
}
// ------------------------------------------------------------------------
@@ -359,13 +364,28 @@ public class HsSubpartitionMemoryDataManager implements HsDataView {
Buffer buffer = bufferConsumer.build();
currentWritingBuffer.close();
bufferConsumer.close();
-
HsBufferContext bufferContext =
- new HsBufferContext(buffer, finishedBufferIndex, targetChannel);
+ new HsBufferContext(
+ compressBuffersIfPossible(buffer), finishedBufferIndex, targetChannel);
addFinishedBuffer(bufferContext);
memoryDataManagerOperation.onBufferFinished();
}
+ private Buffer compressBuffersIfPossible(Buffer buffer) {
+ if (!canBeCompressed(buffer)) {
+ return buffer;
+ }
+ return checkNotNull(bufferCompressor).compressToOriginalBuffer(buffer);
+ }
+
+ /**
+ * Whether the buffer can be compressed or not. Note that event is not compressed because it is
+ * usually small and the size can become even larger after compression.
+ */
+ private boolean canBeCompressed(Buffer buffer) {
+ return bufferCompressor != null && buffer.isBuffer() && buffer.readableBytes() > 0;
+ }
+
@SuppressWarnings("FieldAccessNotGuarded")
// Note that: callWithLock ensure that code block guarded by resultPartitionReadLock and
// subpartitionLock.
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataSpillerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataSpillerTest.java
index c0019a6c960..bcc602fd66d 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataSpillerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsMemoryDataSpillerTest.java
@@ -22,7 +22,6 @@ import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.core.memory.MemorySegment;
import org.apache.flink.core.memory.MemorySegmentFactory;
import org.apache.flink.runtime.io.network.buffer.Buffer;
-import org.apache.flink.runtime.io.network.buffer.BufferCompressor;
import org.apache.flink.runtime.io.network.buffer.FreeingBufferRecycler;
import org.apache.flink.runtime.io.network.buffer.NetworkBuffer;
import org.apache.flink.runtime.io.network.partition.BufferReaderWriterUtil;
@@ -73,21 +72,20 @@ class HsMemoryDataSpillerTest {
}
@ParameterizedTest
- @ValueSource(strings = {"LZ4", "LZO", "ZSTD", "NULL"})
- void testSpillSuccessfully(String compressionFactoryName) throws Exception {
- memoryDataSpiller =
- createMemoryDataSpiller(
- dataFilePath,
- compressionFactoryName.equals("NULL")
- ? null
- : new BufferCompressor(BUFFER_SIZE, compressionFactoryName));
+ @ValueSource(booleans = {false, true})
+ void testSpillSuccessfully(boolean isCompressed) throws Exception {
+ memoryDataSpiller = createMemoryDataSpiller(dataFilePath);
List<BufferWithIdentity> bufferWithIdentityList = new ArrayList<>();
bufferWithIdentityList.addAll(
createBufferWithIdentityList(
- 0, Arrays.asList(Tuple2.of(0, 0), Tuple2.of(1, 1), Tuple2.of(2, 2))));
+ isCompressed,
+ 0,
+ Arrays.asList(Tuple2.of(0, 0), Tuple2.of(1, 1), Tuple2.of(2, 2))));
bufferWithIdentityList.addAll(
createBufferWithIdentityList(
- 0, Arrays.asList(Tuple2.of(4, 0), Tuple2.of(5, 1), Tuple2.of(6, 2))));
+ isCompressed,
+ 0,
+ Arrays.asList(Tuple2.of(4, 0), Tuple2.of(5, 1), Tuple2.of(6, 2))));
CompletableFuture<List<SpilledBuffer>> future =
memoryDataSpiller.spillAsync(bufferWithIdentityList);
List<SpilledBuffer> expectedSpilledBuffers =
@@ -114,6 +112,7 @@ class HsMemoryDataSpillerTest {
.fileOffset);
}));
checkData(
+ isCompressed,
Arrays.asList(
Tuple2.of(0, 0),
Tuple2.of(1, 1),
@@ -129,7 +128,9 @@ class HsMemoryDataSpillerTest {
List<BufferWithIdentity> bufferWithIdentityList = new ArrayList<>();
bufferWithIdentityList.addAll(
createBufferWithIdentityList(
- 0, Arrays.asList(Tuple2.of(0, 0), Tuple2.of(1, 1), Tuple2.of(2, 2))));
+ false,
+ 0,
+ Arrays.asList(Tuple2.of(0, 0), Tuple2.of(1, 1), Tuple2.of(2, 2))));
memoryDataSpiller.close();
assertThatThrownBy(() -> memoryDataSpiller.spillAsync(bufferWithIdentityList))
.isInstanceOf(RejectedExecutionException.class);
@@ -141,12 +142,13 @@ class HsMemoryDataSpillerTest {
List<BufferWithIdentity> bufferWithIdentityList =
new ArrayList<>(
createBufferWithIdentityList(
+ false,
0,
Arrays.asList(Tuple2.of(0, 0), Tuple2.of(1, 1), Tuple2.of(2, 2))));
memoryDataSpiller.spillAsync(bufferWithIdentityList);
// blocked until spill finished.
memoryDataSpiller.release();
- checkData(Arrays.asList(Tuple2.of(0, 0), Tuple2.of(1, 1), Tuple2.of(2, 2)));
+ checkData(false, Arrays.asList(Tuple2.of(0, 0), Tuple2.of(1, 1), Tuple2.of(2, 2)));
}
/**
@@ -156,7 +158,9 @@ class HsMemoryDataSpillerTest {
* @param dataAndIndexes is the list contains pair of (bufferData, bufferIndex).
*/
private static List<BufferWithIdentity> createBufferWithIdentityList(
- int subpartitionId, List<Tuple2<Integer, Integer>> dataAndIndexes) {
+ boolean isCompressed,
+ int subpartitionId,
+ List<Tuple2<Integer, Integer>> dataAndIndexes) {
List<BufferWithIdentity> bufferWithIdentityList = new ArrayList<>();
for (Tuple2<Integer, Integer> dataAndIndex : dataAndIndexes) {
Buffer.DataType dataType =
@@ -169,6 +173,9 @@ class HsMemoryDataSpillerTest {
Buffer buffer =
new NetworkBuffer(
segment, FreeingBufferRecycler.INSTANCE, dataType, BUFFER_SIZE);
+ if (isCompressed) {
+ buffer.setCompressed(true);
+ }
bufferWithIdentityList.add(
new BufferWithIdentity(buffer, dataAndIndex.f1, subpartitionId));
}
@@ -191,7 +198,8 @@ class HsMemoryDataSpillerTest {
return Collections.unmodifiableList(spilledBuffers);
}
- private void checkData(List<Tuple2<Integer, Integer>> dataAndIndexes) throws Exception {
+ private void checkData(boolean isCompressed, List<Tuple2<Integer, Integer>> dataAndIndexes)
+ throws Exception {
FileChannel readChannel = FileChannel.open(dataFilePath, StandardOpenOption.READ);
ByteBuffer headerBuf = BufferReaderWriterUtil.allocatedHeaderBuffer();
MemorySegment segment = MemorySegmentFactory.allocateUnpooledSegment(BUFFER_SIZE);
@@ -200,6 +208,7 @@ class HsMemoryDataSpillerTest {
BufferReaderWriterUtil.readFromByteChannel(
readChannel, headerBuf, segment, (ignore) -> {});
+ assertThat(buffer.isCompressed()).isEqualTo(isCompressed);
assertThat(buffer.readableBytes()).isEqualTo(BUFFER_SIZE);
assertThat(buffer.getNioBufferReadable().order(ByteOrder.nativeOrder()).getInt())
.isEqualTo(dataAndIndex.f0);
@@ -208,11 +217,6 @@ class HsMemoryDataSpillerTest {
}
private static HsMemoryDataSpiller createMemoryDataSpiller(Path dataFilePath) throws Exception {
- return new HsMemoryDataSpiller(dataFilePath, null);
- }
-
- private static HsMemoryDataSpiller createMemoryDataSpiller(
- Path dataFilePath, BufferCompressor bufferCompressor) throws Exception {
- return new HsMemoryDataSpiller(dataFilePath, bufferCompressor);
+ return new HsMemoryDataSpiller(dataFilePath);
}
}
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManagerTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManagerTest.java
index c8b435df3e5..ba53f2a9efb 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManagerTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/partition/hybrid/HsSubpartitionMemoryDataManagerTest.java
@@ -26,12 +26,18 @@ import org.apache.flink.runtime.io.network.api.serialization.EventSerializer;
import org.apache.flink.runtime.io.network.buffer.Buffer;
import org.apache.flink.runtime.io.network.buffer.Buffer.DataType;
import org.apache.flink.runtime.io.network.buffer.BufferBuilder;
+import org.apache.flink.runtime.io.network.buffer.BufferCompressor;
+import org.apache.flink.runtime.io.network.buffer.BufferDecompressor;
import org.apache.flink.runtime.io.network.buffer.ReadOnlySlicedNetworkBuffer;
import org.apache.flink.runtime.io.network.partition.ResultSubpartition.BufferAndBacklog;
import org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.ConsumeStatus;
import org.apache.flink.runtime.io.network.partition.hybrid.HsSpillingInfoProvider.SpillStatus;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+
+import javax.annotation.Nullable;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
@@ -45,10 +51,10 @@ import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import java.util.stream.Collectors;
+import java.util.stream.IntStream;
import static org.apache.flink.runtime.io.network.partition.hybrid.HybridShuffleTestUtils.createBufferBuilder;
import static org.apache.flink.runtime.io.network.partition.hybrid.HybridShuffleTestUtils.createTestingOutputMetrics;
-import static org.apache.flink.util.Preconditions.checkArgument;
import static org.assertj.core.api.Assertions.assertThat;
/** Tests for {@link HsSubpartitionMemoryDataManager}. */
@@ -57,7 +63,7 @@ class HsSubpartitionMemoryDataManagerTest {
private static final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
- private static final int RECORD_SIZE = Integer.BYTES;
+ private static final int RECORD_SIZE = Long.BYTES;
private int bufferSize = RECORD_SIZE;
@@ -202,34 +208,54 @@ class HsSubpartitionMemoryDataManagerTest {
.isInstanceOf(ReadOnlySlicedNetworkBuffer.class)));
}
- @Test
- void testConsumeBuffer() throws Exception {
+ @ParameterizedTest
+ @ValueSource(strings = {"LZ4", "LZO", "ZSTD", "NULL"})
+ void testConsumeBuffer(String compressionFactoryName) throws Exception {
+ final int numDataBuffers = 10;
+ final int numRecordsPerBuffer = 10;
+ // write numRecordsPerBuffer long record to one buffer, as a single long is
+ // incompressible.
+ bufferSize = RECORD_SIZE * numRecordsPerBuffer;
+ BufferCompressor bufferCompressor =
+ compressionFactoryName.equals("NULL")
+ ? null
+ : new BufferCompressor(bufferSize, compressionFactoryName);
+ BufferDecompressor bufferDecompressor =
+ compressionFactoryName.equals("NULL")
+ ? null
+ : new BufferDecompressor(bufferSize, compressionFactoryName);
+
List<BufferIndexAndChannel> consumedBufferIndexAndChannel = new ArrayList<>();
TestingMemoryDataManagerOperation memoryDataManagerOperation =
TestingMemoryDataManagerOperation.builder()
- .setRequestBufferFromPoolSupplier(() -> createBufferBuilder(RECORD_SIZE))
+ .setRequestBufferFromPoolSupplier(() -> createBufferBuilder(bufferSize))
.setOnBufferConsumedConsumer(consumedBufferIndexAndChannel::add)
.build();
HsSubpartitionMemoryDataManager subpartitionMemoryDataManager =
- createSubpartitionMemoryDataManager(memoryDataManagerOperation);
-
- subpartitionMemoryDataManager.append(createRecord(0), DataType.DATA_BUFFER);
- subpartitionMemoryDataManager.append(createRecord(1), DataType.DATA_BUFFER);
- subpartitionMemoryDataManager.append(createRecord(2), DataType.EVENT_BUFFER);
+ createSubpartitionMemoryDataManager(memoryDataManagerOperation, bufferCompressor);
+ List<Tuple2<Long, Buffer.DataType>> expectedRecords = new ArrayList<>();
+
+ long recordValue = 0L;
+ for (int i = 0; i < numDataBuffers; i++) {
+ for (int j = 0; j < numRecordsPerBuffer; j++) {
+ subpartitionMemoryDataManager.append(
+ createRecord(recordValue), DataType.DATA_BUFFER);
+ expectedRecords.add(Tuple2.of(recordValue++, DataType.DATA_BUFFER));
+ }
+ }
+ subpartitionMemoryDataManager.append(createRecord(recordValue), DataType.EVENT_BUFFER);
+ expectedRecords.add(Tuple2.of(recordValue, DataType.EVENT_BUFFER));
- List<Tuple2<Integer, Buffer.DataType>> expectedRecords = new ArrayList<>();
- expectedRecords.add(Tuple2.of(0, Buffer.DataType.DATA_BUFFER));
- expectedRecords.add(Tuple2.of(1, Buffer.DataType.DATA_BUFFER));
- expectedRecords.add(Tuple2.of(2, DataType.EVENT_BUFFER));
+ ArrayList<Optional<BufferAndBacklog>> bufferAndBacklogOpts = new ArrayList<>();
+ for (int i = 0; i < numDataBuffers + 1; i++) {
+ bufferAndBacklogOpts.add(subpartitionMemoryDataManager.consumeBuffer(i));
+ }
checkConsumedBufferAndNextDataType(
- expectedRecords,
- Arrays.asList(
- subpartitionMemoryDataManager.consumeBuffer(0),
- subpartitionMemoryDataManager.consumeBuffer(1),
- subpartitionMemoryDataManager.consumeBuffer(2)));
+ numRecordsPerBuffer, bufferDecompressor, expectedRecords, bufferAndBacklogOpts);
List<BufferIndexAndChannel> expectedBufferIndexAndChannel =
- HybridShuffleTestUtils.createBufferIndexAndChannelsList(0, 0, 1, 2);
+ HybridShuffleTestUtils.createBufferIndexAndChannelsList(
+ 0, IntStream.range(0, numDataBuffers + 1).toArray());
assertThat(consumedBufferIndexAndChannel)
.zipSatisfy(
expectedBufferIndexAndChannel,
@@ -417,29 +443,44 @@ class HsSubpartitionMemoryDataManagerTest {
}
private static void checkConsumedBufferAndNextDataType(
- List<Tuple2<Integer, Buffer.DataType>> expectedRecords,
+ int numRecordsPerBuffer,
+ BufferDecompressor bufferDecompressor,
+ List<Tuple2<Long, Buffer.DataType>> expectedRecords,
List<Optional<BufferAndBacklog>> bufferAndBacklogOpt) {
- checkArgument(expectedRecords.size() == bufferAndBacklogOpt.size());
for (int i = 0; i < bufferAndBacklogOpt.size(); i++) {
- final int index = i;
- assertThat(bufferAndBacklogOpt.get(index))
+ final int bufferIndex = i;
+ assertThat(bufferAndBacklogOpt.get(bufferIndex))
.hasValueSatisfying(
(bufferAndBacklog -> {
Buffer buffer = bufferAndBacklog.buffer();
- int value =
+ if (buffer.isCompressed()) {
+ assertThat(bufferDecompressor).isNotNull();
+ buffer =
+ bufferDecompressor.decompressToIntermediateBuffer(
+ buffer);
+ }
+ ByteBuffer byteBuffer =
buffer.getNioBufferReadable()
- .order(ByteOrder.LITTLE_ENDIAN)
- .getInt();
- Buffer.DataType dataType = buffer.getDataType();
- assertThat(value).isEqualTo(expectedRecords.get(index).f0);
- assertThat(dataType).isEqualTo(expectedRecords.get(index).f1);
- if (index != bufferAndBacklogOpt.size() - 1) {
+ .order(ByteOrder.LITTLE_ENDIAN);
+ int recordIndex = bufferIndex * numRecordsPerBuffer;
+ while (byteBuffer.hasRemaining()) {
+ long value = byteBuffer.getLong();
+ Buffer.DataType dataType = buffer.getDataType();
+ assertThat(value)
+ .isEqualTo(expectedRecords.get(recordIndex).f0);
+ assertThat(dataType)
+ .isEqualTo(expectedRecords.get(recordIndex).f1);
+ recordIndex++;
+ }
+
+ if (bufferIndex != bufferAndBacklogOpt.size() - 1) {
assertThat(bufferAndBacklog.getNextDataType())
- .isEqualTo(expectedRecords.get(index + 1).f1);
+ .isEqualTo(expectedRecords.get(recordIndex).f1);
} else {
assertThat(bufferAndBacklog.getNextDataType())
.isEqualTo(Buffer.DataType.NONE);
}
+ buffer.recycleBuffer();
}));
}
}
@@ -459,17 +500,27 @@ class HsSubpartitionMemoryDataManagerTest {
private HsSubpartitionMemoryDataManager createSubpartitionMemoryDataManager(
HsMemoryDataManagerOperation memoryDataManagerOperation) {
+ return createSubpartitionMemoryDataManager(memoryDataManagerOperation, null);
+ }
+
+ private HsSubpartitionMemoryDataManager createSubpartitionMemoryDataManager(
+ HsMemoryDataManagerOperation memoryDataManagerOperation,
+ @Nullable BufferCompressor bufferCompressor) {
HsSubpartitionMemoryDataManager subpartitionMemoryDataManager =
new HsSubpartitionMemoryDataManager(
- SUBPARTITION_ID, bufferSize, lock.readLock(), memoryDataManagerOperation);
+ SUBPARTITION_ID,
+ bufferSize,
+ lock.readLock(),
+ bufferCompressor,
+ memoryDataManagerOperation);
subpartitionMemoryDataManager.setOutputMetrics(createTestingOutputMetrics());
return subpartitionMemoryDataManager;
}
- private static ByteBuffer createRecord(int value) {
+ private static ByteBuffer createRecord(long value) {
ByteBuffer byteBuffer = ByteBuffer.allocate(RECORD_SIZE);
byteBuffer.order(ByteOrder.LITTLE_ENDIAN);
- byteBuffer.putInt(value);
+ byteBuffer.putLong(value);
byteBuffer.flip();
return byteBuffer;
}