This is an automated email from the ASF dual-hosted git repository.
SteNicholas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new b10546448 [CELEBORN-2348] Support end-to-end shuffle integrity check
for Flink
b10546448 is described below
commit b10546448657d430cc33d8607012a34135217498
Author: Nicholas Jiang <[email protected]>
AuthorDate: Fri Jun 5 10:42:12 2026 +0800
[CELEBORN-2348] Support end-to-end shuffle integrity check for Flink
### What changes were proposed in this pull request?
This PR extends the end-to-end shuffle integrity checks introduced in
CELEBORN-894 (Spark-only) to Flink workloads, covering both the regular and the
tiered (hybrid) read paths. When the check is enabled, the write side records a
per subpartition CRC32 + byte count and the driver validates it against what
the reader actually consumed, failing the read on a mismatch.
- **Write side**: `FlinkShuffleClientImpl` hashes each push payload (the
body after the batch header) into `PushState` via a zero-copy `ByteBuffer` view
and reports the per-subpartition CRC32/bytes at `MapperEnd`, reusing the
existing `crc32PerPartition` / `bytesWrittenPerPartition` plumbing. The
constructor fails fast if the write-side `BATCH_HEADER_SIZE` ever diverges from
the read-side `BufferUtils.HEADER_LENGTH_PREFIX`.
- **Read side**: `RemoteBufferStreamReader` and
`CelebornChannelBufferReader` accumulate the read CRC32/bytes through a shared
`ReadIntegrityTracker` and report them at the last partition's stream end. The
tracker owns the per-path framing/stripping and disables itself on any
unexpected buffer shape (wrong component count, or a buffer shorter than the
batch header) rather than risk a false mismatch.
- **Driver side**: `ReadReducerPartitionEnd` is reused for MAP partitions,
and `MapPartitionCommitHandler.finishPartition` combines the recorded
write-side checksums over the consumed subpartition range, failing closed on a
mismatch or missing metadata.
- Add zero-copy `ByteBuffer` overloads to `CelebornCRC32`, `CommitMetadata`
and `PushState`.
- Minor: drop a stray (cosmetic, no-op) unary plus in
`handleReducerPartitionEnd`'s failure branch and update the client config doc.
### Why are the changes needed?
CELEBORN-894 added end-to-end integrity verification only for Spark. Flink
workloads — including hybrid/tiered shuffle — had no equivalent guard, so
silent shuffle data corruption (bit flips, truncation, mis-framing) could go
undetected and surface as wrong results rather than a failed task. This PR
brings the same write-vs-read checksum/byte-count validation to Flink so such
corruption fails the read instead of being silently consumed.
### Does this PR resolve a correctness bug?
- [ ] Yes
### Does this PR introduce _any_ user-facing change?
- [x] Yes
The existing `celeborn.client.shuffle.integrityCheck.enabled` config now
also applies to Flink (previously Spark-only); its documentation is updated
accordingly. The default remains `false`, so there is no behavior change unless
the check is explicitly enabled.
### How was this patch tested?
Added unit and integration tests:
- `CelebornCRC32Test` / `CommitMetadataTest`: the new `ByteBuffer`
overloads (single and split header/data), order-independence, and corruption /
byte-count-mismatch detection.
- `MapPartitionCommitHandlerTest`: `finishPartition` success and all
failure branches (no metadata, missing map partition, out-of-bounds range,
checksum mismatch, byte-count mismatch), concurrent recording, and the
expired-shuffle race.
- `ReadIntegrityTrackerTest`: report-once / disable semantics and per-path
framing for both the regular and tiered read paths.
- `RemoteBufferStreamReaderTest`: the stream-end-after-close race (a failed
report must not notify the failure listener on a closed channel).
- `CelebornBufferStreamTest`: the `hasRemainingPartitions` location-index
boundary.
- `WordCountTest` (`WordCountTestWithIntegrityCheck`) and
`HybridShuffleWordCountTest`: end-to-end Flink runs with the check enabled, on
both the regular and hybrid shuffle paths.
Closes #3718 from SteNicholas/CELEBORN-2348.
Authored-by: Nicholas Jiang <[email protected]>
Signed-off-by: Nicholas Jiang <[email protected]>
---
.../flink/tiered/CelebornChannelBufferReader.java | 37 ++-
.../flink/tiered/CelebornTierConsumerAgent.java | 3 +-
.../plugin/flink/ReadIntegrityTracker.java | 174 ++++++++++++
.../plugin/flink/RemoteBufferStreamReader.java | 29 +-
.../flink/RemoteShuffleInputGateDelegation.java | 6 +-
.../plugin/flink/client/CelebornBufferStream.java | 16 ++
.../flink/client/FlinkShuffleClientImpl.java | 17 ++
.../plugin/flink/ReadIntegrityTrackerTest.java | 264 ++++++++++++++++++
.../plugin/flink/RemoteBufferStreamReaderTest.java | 90 +++++++
.../flink/client/CelebornBufferStreamTest.java | 51 ++++
.../apache/celeborn/client/ShuffleClientImpl.java | 2 +-
.../apache/celeborn/client/LifecycleManager.scala | 22 +-
.../celeborn/client/commit/CommitHandler.scala | 4 +-
.../client/commit/MapPartitionCommitHandler.scala | 107 +++++++-
.../commit/MapPartitionCommitHandlerTest.scala | 299 +++++++++++++++++++++
.../org/apache/celeborn/common/CelebornCRC32.java | 22 ++
.../org/apache/celeborn/common/CommitMetadata.java | 18 ++
.../apache/celeborn/common/write/PushState.java | 7 +
.../org/apache/celeborn/common/CelebornConf.scala | 2 +-
.../apache/celeborn/common/CelebornCRC32Test.java | 44 +++
.../apache/celeborn/common/CommitMetadataTest.java | 87 ++++++
docs/configuration/client.md | 2 +-
.../tests/flink/HybridShuffleWordCountTest.scala | 22 +-
.../celeborn/tests/flink/WordCountTest.scala | 7 +
24 files changed, 1310 insertions(+), 22 deletions(-)
diff --git
a/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornChannelBufferReader.java
b/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornChannelBufferReader.java
index e033d1283..94659e681 100644
---
a/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornChannelBufferReader.java
+++
b/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornChannelBufferReader.java
@@ -43,10 +43,12 @@ import org.apache.celeborn.common.network.util.NettyUtils;
import org.apache.celeborn.common.protocol.PbNotifyRequiredSegment;
import org.apache.celeborn.common.protocol.PbReadAddCredit;
import org.apache.celeborn.common.util.JavaUtils;
+import org.apache.celeborn.plugin.flink.ReadIntegrityTracker;
import org.apache.celeborn.plugin.flink.ShuffleResourceDescriptor;
import org.apache.celeborn.plugin.flink.client.CelebornBufferStream;
import org.apache.celeborn.plugin.flink.client.FlinkShuffleClientImpl;
import org.apache.celeborn.plugin.flink.protocol.SubPartitionReadData;
+import org.apache.celeborn.plugin.flink.utils.BufferUtils;
/** Wrap the {@link CelebornBufferStream}, used in flink hybrid shuffle
integration strategy now. */
public class CelebornChannelBufferReader {
@@ -80,6 +82,9 @@ public class CelebornChannelBufferReader {
private volatile ConcurrentHashMap<Integer, Integer>
subPartitionRequiredSegmentIds;
+ // Reports the read CRC32/bytes over the data payload at stream end.
+ private final ReadIntegrityTracker integrityTracker;
+
/** Note this field is to record the number of backlog before the read is
set up. */
private int numBackLog = 0;
@@ -90,7 +95,8 @@ public class CelebornChannelBufferReader {
int startSubIdx,
int endSubIdx,
BiConsumer<ByteBuf, TieredStorageSubpartitionId> dataListener,
- BiConsumer<Throwable, TieredStorageSubpartitionId> failureListener) {
+ BiConsumer<Throwable, TieredStorageSubpartitionId> failureListener,
+ boolean integrityCheckEnabled) {
this.client = client;
this.shuffleId = shuffleDescriptor.getShuffleId();
this.partitionId = shuffleDescriptor.getPartitionId();
@@ -99,6 +105,14 @@ public class CelebornChannelBufferReader {
this.subPartitionIndexEnd = endSubIdx;
this.dataListener = dataListener;
this.failureListener = failureListener;
+ this.integrityTracker =
+ new ReadIntegrityTracker(
+ client,
+ this.shuffleId,
+ this.partitionId,
+ subPartitionIndexStart,
+ subPartitionIndexEnd,
+ integrityCheckEnabled);
this.subPartitionRequiredSegmentIds = JavaUtils.newConcurrentHashMap();
for (int subPartitionId = subPartitionIndexStart;
subPartitionId <= subPartitionIndexEnd;
@@ -292,17 +306,18 @@ public class CelebornChannelBufferReader {
}
public void dataReceived(SubPartitionReadData readData) {
+ ByteBuf flinkBuffer = readData.getFlinkBuffer();
LOG.debug(
"Remote buffer stream reader get stream id {} subPartitionId {}
received readable bytes {}.",
readData.getStreamId(),
readData.getSubPartitionId(),
- readData.getFlinkBuffer().readableBytes());
+ flinkBuffer.readableBytes());
checkState(
readData.getSubPartitionId() >= subPartitionIndexStart
&& readData.getSubPartitionId() <= subPartitionIndexEnd,
"Wrong sub partition id: " + readData.getSubPartitionId());
- dataListener.accept(
- readData.getFlinkBuffer(), new
TieredStorageSubpartitionId(readData.getSubPartitionId()));
+ integrityTracker.accumulateTieredBuffer(flinkBuffer,
BufferUtils.HEADER_LENGTH_PREFIX);
+ dataListener.accept(flinkBuffer, new
TieredStorageSubpartitionId(readData.getSubPartitionId()));
int numRequested = bufferManager.tryRequestBuffersIfNeeded();
if (numRequested > 0) {
bufferManager.decreaseRequiredCredits(numRequested);
@@ -314,7 +329,21 @@ public class CelebornChannelBufferReader {
long streamId = streamEnd.getStreamId();
LOG.debug("Buffer stream reader get stream end for {}", streamId);
if (!closed && !CelebornBufferStream.isEmptyStream(bufferStream)) {
+ // Check before the move, which would otherwise advance the location
index.
+ boolean lastPartition = !bufferStream.hasRemainingPartitions();
bufferStream.moveToNextPartitionIfPossible(streamId,
this::sendRequireSegmentId, true);
+ if (lastPartition) {
+ integrityTracker.report(
+ e -> {
+ if (!closed) {
+ for (int subPartitionId = subPartitionIndexStart;
+ subPartitionId <= subPartitionIndexEnd;
+ subPartitionId++) {
+ failureListener.accept(e, new
TieredStorageSubpartitionId(subPartitionId));
+ }
+ }
+ });
+ }
}
}
diff --git
a/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierConsumerAgent.java
b/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierConsumerAgent.java
index 7bfd9d163..04c361dfd 100644
---
a/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierConsumerAgent.java
+++
b/client-flink/common-tiered/src/main/java/org/apache/celeborn/plugin/flink/tiered/CelebornTierConsumerAgent.java
@@ -434,7 +434,8 @@ public class CelebornTierConsumerAgent implements
TierConsumerAgent {
subPartitionIdSet.getStartIndex(),
subPartitionIdSet.getEndIndex(),
getDataListener(partitionId),
- getFailureListener(partitionId));
+ getFailureListener(partitionId),
+ conf.clientShuffleIntegrityCheckEnabled());
for (int id = subPartitionIdSet.getStartIndex(); id <=
subPartitionIdSet.getEndIndex(); id++) {
TieredStorageSubpartitionId subPartitionId = new
TieredStorageSubpartitionId(id);
diff --git
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/ReadIntegrityTracker.java
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/ReadIntegrityTracker.java
new file mode 100644
index 000000000..1b98b52f0
--- /dev/null
+++
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/ReadIntegrityTracker.java
@@ -0,0 +1,174 @@
+/*
+ * 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.plugin.flink;
+
+import java.io.IOException;
+import java.util.function.Consumer;
+
+import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf;
+import org.apache.flink.shaded.netty4.io.netty.buffer.CompositeByteBuf;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.celeborn.common.CommitMetadata;
+import org.apache.celeborn.plugin.flink.client.FlinkShuffleClientImpl;
+import org.apache.celeborn.plugin.flink.utils.BufferUtils;
+
+/**
+ * Accumulates the read-side {@link CommitMetadata} (CRC32 + byte count) over
a consumed
+ * subpartition range and reports it to the driver at stream end. Each read
path strips its own
+ * framing via {@code accumulate*Buffer} so the hashed bytes match what the
writer hashed (the body
+ * after the batch header).
+ */
+public class ReadIntegrityTracker {
+ private static final Logger LOG =
LoggerFactory.getLogger(ReadIntegrityTracker.class);
+
+ private final FlinkShuffleClientImpl client;
+ private final int shuffleId;
+ private final int partitionId;
+ private final int subPartitionIndexStart;
+ private final int subPartitionIndexEnd;
+ private final CommitMetadata readCommitMetadata = new CommitMetadata();
+ private volatile boolean enabled;
+ // True once report() has run; the report is terminal whether it succeeded
or failed.
+ private volatile boolean reportAttempted = false;
+
+ public ReadIntegrityTracker(
+ FlinkShuffleClientImpl client,
+ int shuffleId,
+ int partitionId,
+ int subPartitionIndexStart,
+ int subPartitionIndexEnd,
+ boolean enabled) {
+ this.client = client;
+ this.shuffleId = shuffleId;
+ this.partitionId = partitionId;
+ this.subPartitionIndexStart = subPartitionIndexStart;
+ this.subPartitionIndexEnd = subPartitionIndexEnd;
+ this.enabled = enabled;
+ }
+
+ public boolean isEnabled() {
+ return enabled;
+ }
+
+ /** Permanently disables tracking for this stream so nothing is reported at
stream end. */
+ void disable() {
+ enabled = false;
+ }
+
+ /** Strips the {@code headerPrefix}-byte batch header; a too-short buffer
disables tracking. */
+ private void accumulatePlainBuffer(ByteBuf flinkBuffer, int headerPrefix) {
+ if (flinkBuffer.readableBytes() < headerPrefix) {
+ disableForUnexpectedFraming(
+ String.format("a buffer shorter than the %d-byte batch header",
headerPrefix));
+ return;
+ }
+ readCommitMetadata.addData(
+ flinkBuffer.nioBuffer(
+ flinkBuffer.readerIndex() + headerPrefix,
flinkBuffer.readableBytes() - headerPrefix));
+ }
+
+ /**
+ * Regular path: strips the {@link BufferUtils#HEADER_LENGTH_PREFIX}-byte
header. This path never
+ * splits buffers, so a composite disables tracking.
+ */
+ public void accumulateRegularBuffer(ByteBuf flinkBuffer) {
+ if (!enabled) {
+ return;
+ }
+ if (flinkBuffer instanceof CompositeByteBuf) {
+ disableForUnexpectedFraming("an unexpected composite buffer on the
regular read path");
+ return;
+ }
+ accumulatePlainBuffer(flinkBuffer, BufferUtils.HEADER_LENGTH_PREFIX);
+ }
+
+ /**
+ * Tiered path: strips the {@code headerPrefix}-byte header. Accepts a plain
buffer or a
+ * two-component composite (header + data); other shapes disable tracking.
+ */
+ public void accumulateTieredBuffer(ByteBuf flinkBuffer, int headerPrefix) {
+ if (!enabled) {
+ return;
+ }
+ if (flinkBuffer instanceof CompositeByteBuf) {
+ CompositeByteBuf composite = (CompositeByteBuf) flinkBuffer;
+ if (composite.numComponents() != 2) {
+ disableForUnexpectedFraming(
+ String.format("a composite buffer with %d components",
composite.numComponents()));
+ return;
+ }
+ ByteBuf header = composite.component(0);
+ ByteBuf data = composite.component(1);
+ if (header.readableBytes() < headerPrefix) {
+ disableForUnexpectedFraming(
+ String.format(
+ "a header component shorter than the %d-byte batch header",
headerPrefix));
+ return;
+ }
+ readCommitMetadata.addData(
+ header.nioBuffer(
+ header.readerIndex() + headerPrefix, header.readableBytes() -
headerPrefix),
+ data.nioBuffer(data.readerIndex(), data.readableBytes()));
+ } else {
+ accumulatePlainBuffer(flinkBuffer, headerPrefix);
+ }
+ }
+
+ /**
+ * Reports the accumulated metadata to the driver once, at the last
partition's stream end. A
+ * failed report is logged and passed to {@code onFailure} — terminal, not
retried, since that
+ * already fails the channel.
+ */
+ public void report(Consumer<IOException> onFailure) {
+ if (!enabled || reportAttempted) {
+ return;
+ }
+ reportAttempted = true;
+ try {
+ client.readReducerPartitionEnd(
+ shuffleId,
+ partitionId,
+ subPartitionIndexStart,
+ subPartitionIndexEnd,
+ readCommitMetadata.getChecksum(),
+ readCommitMetadata.getBytes());
+ } catch (IOException e) {
+ LOG.error(
+ "Integrity check report failed for shuffle {} partition {}
subpartitions [{}, {}].",
+ shuffleId,
+ partitionId,
+ subPartitionIndexStart,
+ subPartitionIndexEnd,
+ e);
+ onFailure.accept(e);
+ }
+ }
+
+ private void disableForUnexpectedFraming(String reason) {
+ disable();
+ LOG.warn(
+ "Integrity check skipped for shuffle {} partition {} subpartitions
[{}, {}] because the "
+ + "reader saw {}.",
+ shuffleId,
+ partitionId,
+ subPartitionIndexStart,
+ subPartitionIndexEnd,
+ reason);
+ }
+}
diff --git
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteBufferStreamReader.java
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteBufferStreamReader.java
index 632fa8793..37c18341e 100644
---
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteBufferStreamReader.java
+++
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteBufferStreamReader.java
@@ -50,6 +50,8 @@ public class RemoteBufferStreamReader extends CreditListener {
private CelebornBufferStream bufferStream;
private volatile boolean closed = false;
private Consumer<RequestMessage> messageConsumer;
+ // Reports the read CRC32/bytes at stream end; disabled lazily on unexpected
composite buffers.
+ private final ReadIntegrityTracker integrityTracker;
public RemoteBufferStreamReader(
FlinkShuffleClientImpl client,
@@ -58,7 +60,8 @@ public class RemoteBufferStreamReader extends CreditListener {
int endSubIdx,
TransferBufferPool bufferPool,
Consumer<ByteBuf> dataListener,
- Consumer<Throwable> failureListener) {
+ Consumer<Throwable> failureListener,
+ boolean integrityCheckEnabled) {
this.client = client;
this.shuffleId = shuffleDescriptor.getShuffleId();
this.partitionId = shuffleDescriptor.getPartitionId();
@@ -67,6 +70,14 @@ public class RemoteBufferStreamReader extends CreditListener
{
this.subPartitionIndexEnd = endSubIdx;
this.dataListener = dataListener;
this.failureListener = failureListener;
+ this.integrityTracker =
+ new ReadIntegrityTracker(
+ client,
+ this.shuffleId,
+ this.partitionId,
+ subPartitionIndexStart,
+ subPartitionIndexEnd,
+ integrityCheckEnabled);
this.messageConsumer =
requestMessage -> {
if (requestMessage instanceof ReadData) {
@@ -148,16 +159,28 @@ public class RemoteBufferStreamReader extends
CreditListener {
}
public void dataReceived(ReadData readData) {
+ ByteBuf flinkBuffer = readData.getFlinkBuffer();
logger.debug(
"Remote buffer stream reader get stream id {} received readable bytes
{}.",
readData.getStreamId(),
- readData.getFlinkBuffer().readableBytes());
- dataListener.accept(readData.getFlinkBuffer());
+ flinkBuffer.readableBytes());
+ integrityTracker.accumulateRegularBuffer(flinkBuffer);
+ dataListener.accept(flinkBuffer);
}
public void onStreamEnd(BufferStreamEnd streamEnd) {
long streamId = streamEnd.getStreamId();
logger.debug("Buffer stream reader get stream end for {}", streamId);
+ // Check before the move, which would otherwise advance the location index.
+ boolean lastPartition = !bufferStream.hasRemainingPartitions();
bufferStream.moveToNextPartitionIfPossible(streamId, null, false);
+ if (lastPartition) {
+ integrityTracker.report(
+ e -> {
+ if (!closed) {
+ failureListener.accept(e);
+ }
+ });
+ }
}
}
diff --git
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateDelegation.java
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateDelegation.java
index 99f144c1c..0b496893b 100644
---
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateDelegation.java
+++
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleInputGateDelegation.java
@@ -135,6 +135,8 @@ public class RemoteShuffleInputGateDelegation {
private final ShuffleIOMetricGroup shuffleIOMetricGroup;
+ private final boolean shuffleIntegrityCheckEnabled;
+
public RemoteShuffleInputGateDelegation(
CelebornConf celebornConf,
ShuffleIOOwnerContext ownerContext,
@@ -194,6 +196,7 @@ public class RemoteShuffleInputGateDelegation {
this.numConcurrentReading = numConcurrentReading;
this.availabilityHelper = availabilityHelper;
this.partitionConnectionExceptionEnabled =
celebornConf.partitionConnectionExceptionEnabled();
+ this.shuffleIntegrityCheckEnabled =
celebornConf.clientShuffleIntegrityCheckEnabled();
LOG.debug("Initial input gate with numConcurrentReading {}",
this.numConcurrentReading);
}
@@ -223,7 +226,8 @@ public class RemoteShuffleInputGateDelegation {
endSubIndex,
transferBufferPool,
getDataListener(descriptor.getLeft(), shuffleIOMetricGroup),
- getFailureListener(remoteDescriptor.getResultPartitionID()));
+ getFailureListener(remoteDescriptor.getResultPartitionID()),
+ shuffleIntegrityCheckEnabled);
bufferReaders.add(reader);
numSubPartitionsNotConsumed[descriptor.getLeft()] =
numSubpartitionsPerChannel;
diff --git
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/client/CelebornBufferStream.java
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/client/CelebornBufferStream.java
index 04f8b2387..6588719f2 100644
---
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/client/CelebornBufferStream.java
+++
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/client/CelebornBufferStream.java
@@ -28,6 +28,7 @@ import java.util.function.Supplier;
import javax.annotation.Nullable;
+import com.google.common.annotations.VisibleForTesting;
import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -159,6 +160,21 @@ public class CelebornBufferStream {
return streamId;
}
+ /** Returns whether this stream still has partition locations left to read.
*/
+ public boolean hasRemainingPartitions() {
+ return currentLocationIndex.get() < locations.length;
+ }
+
+ /**
+ * Sets the current location index. {@code openStreamInternal} advances it
via the network stack,
+ * which a unit test cannot drive, so tests poke it directly to reach the
boundary readers rely
+ * on.
+ */
+ @VisibleForTesting
+ void setCurrentLocationIndex(int index) {
+ currentLocationIndex.set(index);
+ }
+
public static CelebornBufferStream create(
FlinkShuffleClientImpl client,
FlinkTransportClientFactory dataClientFactory,
diff --git
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/client/FlinkShuffleClientImpl.java
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/client/FlinkShuffleClientImpl.java
index 9dfe48a9e..81a5cef22 100644
---
a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/client/FlinkShuffleClientImpl.java
+++
b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/client/FlinkShuffleClientImpl.java
@@ -72,6 +72,7 @@ import org.apache.celeborn.common.write.PushState;
import org.apache.celeborn.plugin.flink.network.FlinkTransportClientFactory;
import org.apache.celeborn.plugin.flink.network.ReadClientHandler;
import
org.apache.celeborn.plugin.flink.network.TransportFrameDecoderWithBufferSupplier;
+import org.apache.celeborn.plugin.flink.utils.BufferUtils;
public class FlinkShuffleClientImpl extends ShuffleClientImpl {
public static final Logger logger =
LoggerFactory.getLogger(FlinkShuffleClientImpl.class);
@@ -179,6 +180,16 @@ public class FlinkShuffleClientImpl extends
ShuffleClientImpl {
UserIdentifier userIdentifier,
int bufferSizeBytes) {
super(appUniqueId, conf, userIdentifier);
+ // The integrity check is only correct while the write-side
BATCH_HEADER_SIZE and read-side
+ // BufferUtils.HEADER_LENGTH_PREFIX stay equal, so fail fast if they ever
diverge.
+ if (BATCH_HEADER_SIZE != BufferUtils.HEADER_LENGTH_PREFIX) {
+ throw new IllegalStateException(
+ String.format(
+ "Batch header size mismatch: write-side BATCH_HEADER_SIZE %d !=
read-side "
+ + "BufferUtils.HEADER_LENGTH_PREFIX %d; the Flink shuffle
integrity check would "
+ + "produce false mismatches.",
+ BATCH_HEADER_SIZE, BufferUtils.HEADER_LENGTH_PREFIX));
+ }
this.bufferSizeBytes = bufferSizeBytes;
String module = TransportModuleConstants.DATA_MODULE;
TransportConf dataTransportConf =
@@ -339,6 +350,12 @@ public class FlinkShuffleClientImpl extends
ShuffleClientImpl {
data.writeInt(nextBatchId);
data.writeInt(totalLength - BATCH_HEADER_SIZE);
data.resetWriterIndex();
+ // Hash the payload after the header via a zero-copy ByteBuffer view; the
reader strips the same
+ // count (asserted equal in the constructor).
+ if (shuffleIntegrityCheckEnabled) {
+ int payloadLength = totalLength - BATCH_HEADER_SIZE;
+ pushState.addData(partitionId, data.nioBuffer(BATCH_HEADER_SIZE,
payloadLength));
+ }
logger.debug(
"Do push data byteBuf size {} for app {} shuffle {} map {} attempt {}
reduce {} batch {}.",
totalLength,
diff --git
a/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/ReadIntegrityTrackerTest.java
b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/ReadIntegrityTrackerTest.java
new file mode 100644
index 000000000..d450b44cd
--- /dev/null
+++
b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/ReadIntegrityTrackerTest.java
@@ -0,0 +1,264 @@
+/*
+ * 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.plugin.flink;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertSame;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoInteractions;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.function.Consumer;
+
+import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf;
+import org.apache.flink.shaded.netty4.io.netty.buffer.CompositeByteBuf;
+import org.apache.flink.shaded.netty4.io.netty.buffer.Unpooled;
+import org.junit.Test;
+
+import org.apache.celeborn.common.CommitMetadata;
+import org.apache.celeborn.plugin.flink.client.FlinkShuffleClientImpl;
+import org.apache.celeborn.plugin.flink.utils.BufferUtils;
+
+public class ReadIntegrityTrackerTest {
+
+ private static final int SHUFFLE_ID = 7;
+ private static final int PARTITION_ID = 3;
+ private static final int START_SUB_INDEX = 2;
+ private static final int END_SUB_INDEX = 5;
+ private static final int PREFIX = BufferUtils.HEADER_LENGTH_PREFIX;
+ private static final Consumer<IOException> NO_OP = e -> {};
+
+ private FlinkShuffleClientImpl mockClient() {
+ return mock(FlinkShuffleClientImpl.class);
+ }
+
+ private ReadIntegrityTracker newTracker(FlinkShuffleClientImpl client,
boolean enabled) {
+ return new ReadIntegrityTracker(
+ client, SHUFFLE_ID, PARTITION_ID, START_SUB_INDEX, END_SUB_INDEX,
enabled);
+ }
+
+ /** Deterministic, distinct bytes so different ranges produce different
checksums. */
+ private static byte[] bytesOfLength(int length) {
+ byte[] bytes = new byte[length];
+ for (int i = 0; i < length; i++) {
+ bytes[i] = (byte) (i + 1);
+ }
+ return bytes;
+ }
+
+ /** Prepends a distinct {@link #PREFIX}-byte batch header to {@code
payload}. */
+ private static byte[] frameWithPrefix(byte[] payload) {
+ byte[] framed = new byte[PREFIX + payload.length];
+ for (int i = 0; i < PREFIX; i++) {
+ framed[i] = (byte) (0x80 + i);
+ }
+ System.arraycopy(payload, 0, framed, PREFIX, payload.length);
+ return framed;
+ }
+
+ private void verifyReported(FlinkShuffleClientImpl client, CommitMetadata
expected)
+ throws IOException {
+ verify(client)
+ .readReducerPartitionEnd(
+ SHUFFLE_ID,
+ PARTITION_ID,
+ START_SUB_INDEX,
+ END_SUB_INDEX,
+ expected.getChecksum(),
+ expected.getBytes());
+ }
+
+ @Test
+ public void reportSendsAccumulatedMetadataOnce() throws IOException {
+ FlinkShuffleClientImpl client = mockClient();
+ ReadIntegrityTracker tracker = newTracker(client, true);
+ byte[] payload = "regular-path-payload".getBytes();
+
tracker.accumulateRegularBuffer(Unpooled.wrappedBuffer(frameWithPrefix(payload)));
+
+ tracker.report(NO_OP);
+ tracker.report(NO_OP);
+
+ CommitMetadata expected = new CommitMetadata();
+ expected.addData(ByteBuffer.wrap(payload));
+ verifyReported(client, expected);
+ verify(client, times(1))
+ .readReducerPartitionEnd(anyInt(), anyInt(), anyInt(), anyInt(),
anyInt(), anyLong());
+ }
+
+ @Test
+ public void reportDoesNothingWhenDisabled() {
+ // Disabled at construction (integrity check off).
+ FlinkShuffleClientImpl disabledAtStart = mockClient();
+ newTracker(disabledAtStart, false).report(NO_OP);
+ verifyNoInteractions(disabledAtStart);
+
+ // Disabled mid-stream via disable().
+ FlinkShuffleClientImpl disabledLater = mockClient();
+ ReadIntegrityTracker tracker = newTracker(disabledLater, true);
+ tracker.disable();
+ assertFalse(tracker.isEnabled());
+ tracker.report(NO_OP);
+ verifyNoInteractions(disabledLater);
+ }
+
+ @Test
+ public void reportDoesNothingAfterAccumulateTriggersDisable() {
+ FlinkShuffleClientImpl client = mockClient();
+ ReadIntegrityTracker tracker = newTracker(client, true);
+ // Accumulate one good buffer, then a malformed one that disables tracking
mid-stream.
+ tracker.accumulateTieredBuffer(Unpooled.wrappedBuffer(bytesOfLength(PREFIX
+ 8)), PREFIX);
+ assertTrue(tracker.isEnabled());
+ tracker.accumulateTieredBuffer(Unpooled.wrappedBuffer(bytesOfLength(PREFIX
- 1)), PREFIX);
+ assertFalse(tracker.isEnabled());
+ // report must not send the partial metadata.
+ tracker.report(NO_OP);
+ verifyNoInteractions(client);
+ }
+
+ @Test
+ public void reportPassesFailureToCallbackAndDoesNotRetry() throws
IOException {
+ FlinkShuffleClientImpl client = mockClient();
+ IOException boom = new IOException("report failed");
+ doThrow(boom)
+ .when(client)
+ .readReducerPartitionEnd(anyInt(), anyInt(), anyInt(), anyInt(),
anyInt(), anyLong());
+ ReadIntegrityTracker tracker = newTracker(client, true);
+
+ List<IOException> failures = new ArrayList<>();
+ tracker.report(failures::add);
+ // A failed report is still terminal: a second call must not re-issue the
RPC.
+ tracker.report(failures::add);
+
+ assertEquals(1, failures.size());
+ assertSame(boom, failures.get(0));
+ verify(client, times(1))
+ .readReducerPartitionEnd(anyInt(), anyInt(), anyInt(), anyInt(),
anyInt(), anyLong());
+ }
+
+ @Test
+ public void accumulateRegularBufferStripsBatchHeaderPrefix() throws
IOException {
+ FlinkShuffleClientImpl client = mockClient();
+ ReadIntegrityTracker tracker = newTracker(client, true);
+ // Only the payload after the PREFIX-byte batch header (what the writer
hashed) is accumulated.
+ byte[] framed = bytesOfLength(PREFIX + 24);
+ tracker.accumulateRegularBuffer(Unpooled.wrappedBuffer(framed));
+ assertTrue(tracker.isEnabled());
+ tracker.report(NO_OP);
+
+ CommitMetadata expected = new CommitMetadata();
+ expected.addData(ByteBuffer.wrap(Arrays.copyOfRange(framed, PREFIX,
framed.length)));
+ verifyReported(client, expected);
+ }
+
+ @Test
+ public void accumulateRegularBufferDisablesOnUnexpectedFraming() {
+ // The regular path never splits large buffers, so any composite buffer is
unexpected.
+ CompositeByteBuf composite = Unpooled.compositeBuffer();
+ composite.addComponent(true, Unpooled.wrappedBuffer(bytesOfLength(32)));
+ assertRegularBufferDisables(composite);
+
+ // A plain buffer too short to hold the batch header.
+ assertRegularBufferDisables(Unpooled.wrappedBuffer(bytesOfLength(PREFIX -
1)));
+ }
+
+ /** Feeding {@code bad} to the regular path must disable tracking and report
nothing. */
+ private void assertRegularBufferDisables(ByteBuf bad) {
+ FlinkShuffleClientImpl client = mockClient();
+ ReadIntegrityTracker tracker = newTracker(client, true);
+ tracker.accumulateRegularBuffer(bad);
+ assertFalse(tracker.isEnabled());
+ tracker.report(NO_OP);
+ verifyNoInteractions(client);
+ }
+
+ @Test
+ public void accumulateTieredBufferStripsPrefixFromPlainBuffer() throws
IOException {
+ FlinkShuffleClientImpl client = mockClient();
+ ReadIntegrityTracker tracker = newTracker(client, true);
+ byte[] framed = bytesOfLength(PREFIX + 24);
+ tracker.accumulateTieredBuffer(Unpooled.wrappedBuffer(framed), PREFIX);
+ assertTrue(tracker.isEnabled());
+ tracker.report(NO_OP);
+
+ CommitMetadata expected = new CommitMetadata();
+ expected.addData(ByteBuffer.wrap(Arrays.copyOfRange(framed, PREFIX,
framed.length)));
+ verifyReported(client, expected);
+ }
+
+ @Test
+ public void accumulateTieredBufferHashesHeaderTailAndDataOfComposite()
throws IOException {
+ FlinkShuffleClientImpl client = mockClient();
+ ReadIntegrityTracker tracker = newTracker(client, true);
+ // Large buffers arrive as a header component (batch header) plus a data
component.
+ byte[] header = bytesOfLength(BufferUtils.HEADER_LENGTH);
+ byte[] data = "tiered-data-component".getBytes();
+ CompositeByteBuf composite = Unpooled.compositeBuffer();
+ composite.addComponent(true, Unpooled.wrappedBuffer(header));
+ composite.addComponent(true, Unpooled.wrappedBuffer(data));
+
+ tracker.accumulateTieredBuffer(composite, PREFIX);
+ assertTrue(tracker.isEnabled());
+ tracker.report(NO_OP);
+
+ CommitMetadata expected = new CommitMetadata();
+ expected.addData(
+ ByteBuffer.wrap(Arrays.copyOfRange(header, PREFIX, header.length)),
ByteBuffer.wrap(data));
+ verifyReported(client, expected);
+ }
+
+ @Test
+ public void accumulateTieredBufferDisablesOnUnexpectedFraming() {
+ // A composite with a component count other than the expected header +
data.
+ CompositeByteBuf threeComponents = Unpooled.compositeBuffer();
+ threeComponents.addComponent(
+ true,
Unpooled.wrappedBuffer(bytesOfLength(BufferUtils.HEADER_LENGTH)));
+ threeComponents.addComponent(true,
Unpooled.wrappedBuffer(bytesOfLength(8)));
+ threeComponents.addComponent(true,
Unpooled.wrappedBuffer(bytesOfLength(8)));
+ assertTieredBufferDisables(threeComponents);
+
+ // A composite whose header component is shorter than the batch-header
prefix.
+ CompositeByteBuf shortHeader = Unpooled.compositeBuffer();
+ shortHeader.addComponent(true, Unpooled.wrappedBuffer(bytesOfLength(PREFIX
- 1)));
+ shortHeader.addComponent(true, Unpooled.wrappedBuffer(bytesOfLength(8)));
+ assertTieredBufferDisables(shortHeader);
+
+ // A plain buffer shorter than the batch-header prefix.
+ assertTieredBufferDisables(Unpooled.wrappedBuffer(bytesOfLength(PREFIX -
1)));
+ }
+
+ /** Feeding {@code bad} to the tiered path must disable tracking and report
nothing. */
+ private void assertTieredBufferDisables(ByteBuf bad) {
+ FlinkShuffleClientImpl client = mockClient();
+ ReadIntegrityTracker tracker = newTracker(client, true);
+ tracker.accumulateTieredBuffer(bad, PREFIX);
+ assertFalse(tracker.isEnabled());
+ tracker.report(NO_OP);
+ verifyNoInteractions(client);
+ }
+}
diff --git
a/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/RemoteBufferStreamReaderTest.java
b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/RemoteBufferStreamReaderTest.java
new file mode 100644
index 000000000..bca319ef1
--- /dev/null
+++
b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/RemoteBufferStreamReaderTest.java
@@ -0,0 +1,90 @@
+/*
+ * 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.plugin.flink;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+
+import org.junit.Test;
+
+import org.apache.celeborn.common.network.protocol.BufferStreamEnd;
+import org.apache.celeborn.plugin.flink.client.CelebornBufferStream;
+import org.apache.celeborn.plugin.flink.client.FlinkShuffleClientImpl;
+
+public class RemoteBufferStreamReaderTest {
+
+ private static final int SHUFFLE_ID = 1;
+ private static final int PARTITION_ID = 2;
+ private static final int START_SUB = 0;
+ private static final int END_SUB = 3;
+ private static final long STREAM_ID = 99L;
+
+ /** An opened reader whose last-partition stream end reports and the report
RPC throws. */
+ private RemoteBufferStreamReader openedReaderWithFailingReport(
+ FlinkShuffleClientImpl client, List<Throwable> failures) throws
IOException {
+ CelebornBufferStream stream = mock(CelebornBufferStream.class);
+ when(stream.hasRemainingPartitions()).thenReturn(false);
+ when(client.readBufferedPartition(anyInt(), anyInt(), anyInt(), anyInt(),
anyBoolean()))
+ .thenReturn(stream);
+ doThrow(new IOException("report failed"))
+ .when(client)
+ .readReducerPartitionEnd(anyInt(), anyInt(), anyInt(), anyInt(),
anyInt(), anyLong());
+ ShuffleResourceDescriptor descriptor =
+ new ShuffleResourceDescriptor(SHUFFLE_ID, 0, 0, PARTITION_ID);
+ RemoteBufferStreamReader reader =
+ new RemoteBufferStreamReader(
+ client, descriptor, START_SUB, END_SUB, null, buffer -> {},
failures::add, true);
+ reader.open(0);
+ return reader;
+ }
+
+ @Test
+ public void onStreamEndNotifiesFailureListenerWhenOpen() throws IOException {
+ FlinkShuffleClientImpl client = mock(FlinkShuffleClientImpl.class);
+ List<Throwable> failures = new ArrayList<>();
+ RemoteBufferStreamReader reader = openedReaderWithFailingReport(client,
failures);
+
+ reader.onStreamEnd(new BufferStreamEnd(STREAM_ID));
+
+ assertEquals(1, failures.size());
+ assertTrue(failures.get(0) instanceof IOException);
+ }
+
+ @Test
+ public void onStreamEndSkipsFailureListenerAfterClose() throws IOException {
+ FlinkShuffleClientImpl client = mock(FlinkShuffleClientImpl.class);
+ List<Throwable> failures = new ArrayList<>();
+ RemoteBufferStreamReader reader = openedReaderWithFailingReport(client,
failures);
+ // A close racing the stream end must swallow the report failure, not fail
a dead channel.
+ reader.close();
+
+ reader.onStreamEnd(new BufferStreamEnd(STREAM_ID));
+
+ assertTrue(failures.isEmpty());
+ }
+}
diff --git
a/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/client/CelebornBufferStreamTest.java
b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/client/CelebornBufferStreamTest.java
new file mode 100644
index 000000000..1546d251f
--- /dev/null
+++
b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/client/CelebornBufferStreamTest.java
@@ -0,0 +1,51 @@
+/*
+ * 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.plugin.flink.client;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import org.junit.Test;
+
+import org.apache.celeborn.common.protocol.PartitionLocation;
+
+public class CelebornBufferStreamTest {
+
+ // hasRemainingPartitions only reads locations.length, so null elements
suffice here.
+ private static CelebornBufferStream streamWithLocations(int numLocations) {
+ return new CelebornBufferStream(
+ null, null, "shuffleKey", new PartitionLocation[numLocations], 0, 0,
0L, null);
+ }
+
+ @Test
+ public void hasRemainingPartitionsTracksLocationIndexBoundary() {
+ // openStreamInternal advances currentLocationIndex past the partition it
opened, so the last
+ // partition's stream end sees index == length, where lastPartition =
!hasRemainingPartitions().
+ CelebornBufferStream stream = streamWithLocations(2);
+ assertTrue(stream.hasRemainingPartitions()); // index 0 < 2
+ stream.setCurrentLocationIndex(1);
+ assertTrue(stream.hasRemainingPartitions()); // index 1 < 2, more
partitions remain
+ stream.setCurrentLocationIndex(2);
+ assertFalse(stream.hasRemainingPartitions()); // index 2 == 2, last
partition consumed
+ }
+
+ @Test
+ public void hasRemainingPartitionsIsFalseWithoutLocations() {
+ assertFalse(streamWithLocations(0).hasRemainingPartitions());
+ }
+}
diff --git
a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
index be2bdf87d..225a9cb99 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
@@ -120,7 +120,7 @@ public class ShuffleClientImpl extends ShuffleClient {
private final boolean pushExcludeWorkerOnFailureEnabled;
private final boolean shuffleCompressionEnabled;
- private final boolean shuffleIntegrityCheckEnabled;
+ protected final boolean shuffleIntegrityCheckEnabled;
private final Set<String> pushExcludedWorkers =
ConcurrentHashMap.newKeySet();
private final ConcurrentHashMap<String, Long> fetchExcludedWorkers =
diff --git
a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
index d8c48420d..f9508cfe6 100644
--- a/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala
@@ -466,6 +466,9 @@ class LifecycleManager(val appUniqueId: String, val conf:
CelebornConf) extends
attemptId,
partitionId,
numMappers,
+ numPartitions,
+ crc32PerPartition,
+ bytesWrittenPerPartition,
serdeVersion)
case _ =>
throw new UnsupportedOperationException(s"Not support $partitionType
yet")
@@ -474,7 +477,9 @@ class LifecycleManager(val appUniqueId: String, val conf:
CelebornConf) extends
case pb: ReadReducerPartitionEnd =>
val partitionType = getPartitionType(pb.shuffleId)
partitionType match {
- case PartitionType.REDUCE =>
+ // Map partitions reuse this reducer-named RPC/handler; for MAP,
partitionId is the map
+ // partition id and [startMapIndex, endMapIndex] is the consumed
subpartition range.
+ case PartitionType.REDUCE | PartitionType.MAP =>
handleReducerPartitionEnd(
context,
pb.shuffleId,
@@ -559,7 +564,7 @@ class LifecycleManager(val appUniqueId: String, val conf:
CelebornConf) extends
StatusCode.SUCCESS.getValue).build()
} else {
response = PbReadReducerPartitionEndResponse.newBuilder().setStatus(
-
+StatusCode.READ_REDUCER_PARTITION_END_FAILED.getValue).setErrorMsg(errorMessage).build()
+
StatusCode.READ_REDUCER_PARTITION_END_FAILED.getValue).setErrorMsg(errorMessage).build()
}
context.reply(response)
}
@@ -1227,6 +1232,9 @@ class LifecycleManager(val appUniqueId: String, val conf:
CelebornConf) extends
attemptId: Int,
partitionId: Int,
numMappers: Int,
+ numPartitions: Int,
+ crc32PerPartition: Array[Int],
+ bytesWrittenPerPartition: Array[Long],
serdeVersion: SerdeVersion): Unit = {
def reply(result: Boolean): Unit = {
val message =
@@ -1243,7 +1251,15 @@ class LifecycleManager(val appUniqueId: String, val
conf: CelebornConf) extends
}
val (mapperAttemptFinishedSuccess, _) =
- commitManager.finishMapperAttempt(shuffleId, mapId, attemptId,
numMappers, partitionId)
+ commitManager.finishMapperAttempt(
+ shuffleId,
+ mapId,
+ attemptId,
+ numMappers,
+ partitionId,
+ numPartitions = numPartitions,
+ crc32PerPartition = crc32PerPartition,
+ bytesWrittenPerPartition = bytesWrittenPerPartition)
reply(mapperAttemptFinishedSuccess)
}
diff --git
a/client/src/main/scala/org/apache/celeborn/client/commit/CommitHandler.scala
b/client/src/main/scala/org/apache/celeborn/client/commit/CommitHandler.scala
index 4189ee5f7..f54c2b990 100644
---
a/client/src/main/scala/org/apache/celeborn/client/commit/CommitHandler.scala
+++
b/client/src/main/scala/org/apache/celeborn/client/commit/CommitHandler.scala
@@ -432,7 +432,9 @@ abstract class CommitHandler(
}
/**
- * Invoked when a reduce partition finishes reading data to perform end to
end integrity check validation
+ * Invoked when a partition (reduce or map) finishes reading data to perform
end to end integrity
+ * check validation. See `MapPartitionCommitHandler.finishPartition` for the
map-partition
+ * semantics of the parameters.
*/
def finishPartition(
shuffleId: Int,
diff --git
a/client/src/main/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandler.scala
b/client/src/main/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandler.scala
index 9299768b6..04701f2f4 100644
---
a/client/src/main/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandler.scala
+++
b/client/src/main/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandler.scala
@@ -76,6 +76,19 @@ class MapPartitionCommitHandler(
// shuffleId -> boolean, records whether the shuffle is visible at the
segment level, facilitating future optimization of worker read and write
processes
private val shuffleIsSegmentGranularityVisible =
JavaUtils.newConcurrentHashMap[Int, Boolean]
+ private val shuffleIntegrityCheckEnabled =
conf.clientShuffleIntegrityCheckEnabled
+
+ // Write-side per-subpartition checksums of one finished map partition
(indexed by subpartition).
+ private case class MapPartitionWriteMetadata(crc32: Array[Int],
bytesWritten: Array[Long]) {
+ require(
+ crc32.length == bytesWritten.length,
+ s"crc32 length ${crc32.length} != bytesWritten length
${bytesWritten.length}")
+ }
+
+ // shuffleId -> (mapPartitionId -> write-side metadata).
+ private val commitMetadataForMapPartition =
+ JavaUtils.newConcurrentHashMap[Int, ConcurrentHashMap[Int,
MapPartitionWriteMetadata]]()
+
override def getPartitionType(): PartitionType = {
PartitionType.MAP
}
@@ -133,6 +146,7 @@ class MapPartitionCommitHandler(
inProcessMapPartitionEndIds.remove(shuffleId)
shuffleSucceedPartitionIds.remove(shuffleId)
shuffleIsSegmentGranularityVisible.remove(shuffleId)
+ commitMetadataForMapPartition.remove(shuffleId)
super.removeExpiredShuffle(shuffleId)
}
@@ -231,11 +245,46 @@ class MapPartitionCommitHandler(
shuffleId,
(k: Int) => ConcurrentHashMap.newKeySet[Integer]())
resultPartitions.add(partitionId)
+
+ if (shuffleIntegrityCheckEnabled) {
+ recordMapPartitionCommitMetadata(
+ shuffleId,
+ partitionId,
+ numPartitions,
+ crc32PerPartition,
+ bytesWrittenPerPartition)
+ }
}
(dataCommitSuccess, false)
}
+ /**
+ * Records a finished mapper's per-subpartition write-side checksums, keyed
by (shuffleId,
+ * mapPartitionId), for later validation in [[finishPartition]]. A
retried/duplicate attempt
+ * overwrites the prior record (last write wins), which is intentional: the
reader validates
+ * against the last committed attempt. Visible for testing.
+ */
+ private[commit] def recordMapPartitionCommitMetadata(
+ shuffleId: Int,
+ mapPartitionId: Int,
+ numPartitions: Int,
+ crc32PerPartition: Array[Int],
+ bytesWrittenPerPartition: Array[Long]): Unit = {
+ if (crc32PerPartition == null || crc32PerPartition.length != numPartitions
||
+ bytesWrittenPerPartition == null || bytesWrittenPerPartition.length !=
numPartitions) {
+ logWarning(
+ s"Skip recording commit metadata for shuffle $shuffleId map partition
$mapPartitionId " +
+ s"because reported checksum arrays do not match numPartitions
$numPartitions.")
+ return
+ }
+ commitMetadataForMapPartition
+ .computeIfAbsent(
+ shuffleId,
+ (_: Int) => JavaUtils.newConcurrentHashMap[Int,
MapPartitionWriteMetadata]())
+ .put(mapPartitionId, MapPartitionWriteMetadata(crc32PerPartition,
bytesWrittenPerPartition))
+ }
+
override def registerShuffle(
shuffleId: Int,
numMappers: Int,
@@ -243,19 +292,69 @@ class MapPartitionCommitHandler(
numPartitions: Int): Unit = {
super.registerShuffle(shuffleId, numMappers, isSegmentGranularityVisible,
numPartitions)
shuffleIsSegmentGranularityVisible.put(shuffleId,
isSegmentGranularityVisible)
+ // The outer metadata map is created lazily by
recordMapPartitionCommitMetadata.
}
override def isSegmentGranularityVisible(shuffleId: Int): Boolean = {
shuffleIsSegmentGranularityVisible.get(shuffleId)
}
+ /**
+ * Validates a map partition read over the consumed `[startSubIndex,
endSubIndex]` range against
+ * the order-independent combination of the write-side checksums over that
range. The params
+ * rename the trait's reducer-oriented signature
(partitionId/startMapIndex/endMapIndex) to map
+ * semantics. A reader reaches stream end only after commit, so the
write-side metadata is already
+ * recorded; missing metadata fails closed.
+ */
override def finishPartition(
shuffleId: Int,
- partitionId: Int,
- startMapIndex: Int,
- endMapIndex: Int,
+ mapPartitionId: Int,
+ startSubIndex: Int,
+ endSubIndex: Int,
actualCommitMetadata: CommitMetadata): (Boolean, String) = {
- throw new UnsupportedOperationException()
+ if (!shuffleIntegrityCheckEnabled) {
+ return (true, "")
+ }
+ val perMapPartition = commitMetadataForMapPartition.get(shuffleId)
+ if (perMapPartition == null) {
+ return (
+ false,
+ s"No write-side commit metadata recorded for shuffle $shuffleId when
validating " +
+ s"map partition $mapPartitionId subpartitions [$startSubIndex,
$endSubIndex].")
+ }
+ val subPartitionMetadata = perMapPartition.get(mapPartitionId)
+ if (subPartitionMetadata == null) {
+ return (
+ false,
+ s"No write-side commit metadata recorded for shuffle $shuffleId map
partition " +
+ s"$mapPartitionId when validating subpartitions [$startSubIndex,
$endSubIndex].")
+ }
+ val crc32PerSubPartition = subPartitionMetadata.crc32
+ val bytesPerSubPartition = subPartitionMetadata.bytesWritten
+ if (startSubIndex < 0 || endSubIndex >= crc32PerSubPartition.length ||
+ startSubIndex > endSubIndex) {
+ return (
+ false,
+ s"Invalid subpartition range [$startSubIndex, $endSubIndex] for
shuffle $shuffleId " +
+ s"map partition $mapPartitionId with ${crc32PerSubPartition.length}
subpartitions.")
+ }
+ val expectedCommitMetadata = new CommitMetadata()
+ var subIndex = startSubIndex
+ while (subIndex <= endSubIndex) {
+ expectedCommitMetadata.addCommitData(
+ crc32PerSubPartition(subIndex),
+ bytesPerSubPartition(subIndex))
+ subIndex += 1
+ }
+ if (CommitMetadata.checkCommitMetadata(expectedCommitMetadata,
actualCommitMetadata)) {
+ (true, "")
+ } else {
+ (
+ false,
+ s"Integrity check failed for shuffle $shuffleId map partition
$mapPartitionId " +
+ s"subpartitions [$startSubIndex, $endSubIndex], expected
$expectedCommitMetadata " +
+ s"but read $actualCommitMetadata.")
+ }
}
override def handleGetReducerFileGroup(
diff --git
a/client/src/test/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandlerTest.scala
b/client/src/test/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandlerTest.scala
new file mode 100644
index 000000000..47a485b60
--- /dev/null
+++
b/client/src/test/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandlerTest.scala
@@ -0,0 +1,299 @@
+/*
+ * 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.client.commit
+
+import java.util.concurrent.{ConcurrentHashMap, ConcurrentLinkedQueue,
CountDownLatch, ScheduledExecutorService, ThreadPoolExecutor, TimeUnit}
+
+import org.apache.celeborn.CelebornFunSuite
+import org.apache.celeborn.client.{ShuffleCommittedInfo, WorkerStatusTracker}
+import org.apache.celeborn.client.CommitManager.CommittedPartitionInfo
+import org.apache.celeborn.client.LifecycleManager.ShuffleAllocatedWorkers
+import org.apache.celeborn.common.{CelebornConf, CommitMetadata}
+import org.apache.celeborn.common.util.ThreadUtils
+
+class MapPartitionCommitHandlerTest extends CelebornFunSuite {
+
+ // The handler spins up daemon pools; skip the thread audit to avoid flaky
leak warnings.
+ override protected val enableAutoThreadAudit = false
+
+ private var rpcPool: ThreadPoolExecutor = _
+ private var commitScheduler: ScheduledExecutorService = _
+
+ override def beforeAll(): Unit = {
+ super.beforeAll()
+ rpcPool = ThreadUtils.newDaemonCachedThreadPool("test-map-commit-rpc")
+ commitScheduler =
+
ThreadUtils.newDaemonSingleThreadScheduledExecutor("test-map-commit-scheduler")
+ }
+
+ override def afterAll(): Unit = {
+ if (rpcPool != null) {
+ rpcPool.shutdownNow()
+ }
+ if (commitScheduler != null) {
+ commitScheduler.shutdownNow()
+ }
+ super.afterAll()
+ }
+
+ private def newHandler(integrityCheckEnabled: Boolean):
MapPartitionCommitHandler = {
+ val conf = new CelebornConf()
+ conf.set(
+ CelebornConf.CLIENT_SHUFFLE_INTEGRITY_CHECK_ENABLED.key,
+ integrityCheckEnabled.toString)
+ new MapPartitionCommitHandler(
+ "test-app",
+ conf,
+ new ShuffleAllocatedWorkers(),
+ new CommittedPartitionInfo(),
+ new WorkerStatusTracker(conf, null),
+ rpcPool,
+ commitScheduler)
+ }
+
+ /** Write-side per-subpartition checksum/byte arrays computed over the given
payloads. */
+ private def writeSideMetadata(payloads: Seq[Array[Byte]]): (Array[Int],
Array[Long]) = {
+ val crc32PerSubPartition = new Array[Int](payloads.length)
+ val bytesPerSubPartition = new Array[Long](payloads.length)
+ payloads.zipWithIndex.foreach { case (payload, subIndex) =>
+ val commitMetadata = new CommitMetadata()
+ commitMetadata.addDataWithOffsetAndLength(payload, 0, payload.length)
+ crc32PerSubPartition(subIndex) = commitMetadata.getChecksum
+ bytesPerSubPartition(subIndex) = commitMetadata.getBytes
+ }
+ (crc32PerSubPartition, bytesPerSubPartition)
+ }
+
+ /** Read-side commit metadata accumulated over the inclusive subpartition
range. */
+ private def readSideMetadata(
+ payloads: Seq[Array[Byte]],
+ startSubIndex: Int,
+ endSubIndex: Int): CommitMetadata = {
+ val commitMetadata = new CommitMetadata()
+ (startSubIndex to endSubIndex).foreach { subIndex =>
+ val payload = payloads(subIndex)
+ commitMetadata.addDataWithOffsetAndLength(payload, 0, payload.length)
+ }
+ commitMetadata
+ }
+
+ private val shuffleId = 1
+ private val mapPartitionId = 0
+ private val payloads =
+ Seq("sub-0-bytes".getBytes(), "sub-1-bytes".getBytes(),
"sub-2-bytes".getBytes())
+ private val numPartitions = payloads.length
+
+ test("finishPartition returns success without validation when integrity
check is disabled") {
+ val handler = newHandler(integrityCheckEnabled = false)
+ // No metadata recorded, yet a disabled check must short-circuit to
success.
+ val (isValid, message) =
+ handler.finishPartition(shuffleId, mapPartitionId, 0, numPartitions - 1,
new CommitMetadata())
+ assert(isValid)
+ assert(message.isEmpty)
+ }
+
+ test("finishPartition fails when no commit metadata was recorded for the
shuffle") {
+ val handler = newHandler(integrityCheckEnabled = true)
+ // Neither registerShuffle nor recordMapPartitionCommitMetadata called for
this shuffle.
+ val (isValid, message) =
+ handler.finishPartition(shuffleId, mapPartitionId, 0, numPartitions - 1,
new CommitMetadata())
+ assert(!isValid)
+ assert(message.contains(s"No write-side commit metadata recorded for
shuffle $shuffleId"))
+ }
+
+ test("finishPartition fails when no commit metadata was recorded for the map
partition") {
+ val handler = newHandler(integrityCheckEnabled = true)
+ val (crc32PerSubPartition, bytesPerSubPartition) =
writeSideMetadata(payloads)
+ // Record a different map partition so the outer map exists but this one's
entry is absent.
+ handler.recordMapPartitionCommitMetadata(
+ shuffleId,
+ mapPartitionId + 1,
+ numPartitions,
+ crc32PerSubPartition,
+ bytesPerSubPartition)
+ val (isValid, message) =
+ handler.finishPartition(shuffleId, mapPartitionId, 0, numPartitions - 1,
new CommitMetadata())
+ assert(!isValid)
+ assert(message.contains(s"shuffle $shuffleId map partition
$mapPartitionId"))
+ }
+
+ test("recordMapPartitionCommitMetadata skips reported arrays that do not
match numPartitions") {
+ val handler = newHandler(integrityCheckEnabled = true)
+ handler.registerShuffle(shuffleId, 1, isSegmentGranularityVisible = true,
numPartitions)
+ // A crc32 array shorter than numPartitions is rejected, so validation
later fails closed.
+ val (crc32PerSubPartition, bytesPerSubPartition) =
writeSideMetadata(payloads)
+ handler.recordMapPartitionCommitMetadata(
+ shuffleId,
+ mapPartitionId,
+ numPartitions,
+ crc32PerSubPartition.dropRight(1),
+ bytesPerSubPartition)
+ val (isValid, message) =
+ handler.finishPartition(shuffleId, mapPartitionId, 0, numPartitions - 1,
new CommitMetadata())
+ assert(!isValid)
+ assert(message.contains(s"map partition $mapPartitionId"))
+ }
+
+ test("finishPartition fails closed after the shuffle metadata was
expired/removed") {
+ val handler = newHandler(integrityCheckEnabled = true)
+ handler.registerShuffle(shuffleId, 1, isSegmentGranularityVisible = true,
numPartitions)
+ val (crc32PerSubPartition, bytesPerSubPartition) =
writeSideMetadata(payloads)
+ handler.recordMapPartitionCommitMetadata(
+ shuffleId,
+ mapPartitionId,
+ numPartitions,
+ crc32PerSubPartition,
+ bytesPerSubPartition)
+ // Cleanup racing a late reader stream end drops the metadata; validation
must fail, not pass.
+ handler.removeExpiredShuffle(shuffleId)
+ val actual = readSideMetadata(payloads, 0, numPartitions - 1)
+ val (isValid, message) =
+ handler.finishPartition(shuffleId, mapPartitionId, 0, numPartitions - 1,
actual)
+ assert(!isValid)
+ assert(message.contains(s"No write-side commit metadata recorded for
shuffle $shuffleId"))
+ }
+
+ test("finishPartition fails for an out-of-bounds subpartition range") {
+ val handler = newHandler(integrityCheckEnabled = true)
+ val (crc32PerSubPartition, bytesPerSubPartition) =
writeSideMetadata(payloads)
+ handler.recordMapPartitionCommitMetadata(
+ shuffleId,
+ mapPartitionId,
+ numPartitions,
+ crc32PerSubPartition,
+ bytesPerSubPartition)
+ val (isValid, message) =
+ handler.finishPartition(shuffleId, mapPartitionId, 0, numPartitions, new
CommitMetadata())
+ assert(!isValid)
+ assert(message.contains("Invalid subpartition range"))
+ }
+
+ test("finishPartition succeeds when the read-side metadata matches the
write-side range") {
+ val handler = newHandler(integrityCheckEnabled = true)
+ val (crc32PerSubPartition, bytesPerSubPartition) =
writeSideMetadata(payloads)
+ handler.recordMapPartitionCommitMetadata(
+ shuffleId,
+ mapPartitionId,
+ numPartitions,
+ crc32PerSubPartition,
+ bytesPerSubPartition)
+ val startSubIndex = 0
+ val endSubIndex = 1
+ val actual = readSideMetadata(payloads, startSubIndex, endSubIndex)
+ val (isValid, message) =
+ handler.finishPartition(shuffleId, mapPartitionId, startSubIndex,
endSubIndex, actual)
+ assert(isValid)
+ assert(message.isEmpty)
+ }
+
+ test("finishPartition fails when the read-side metadata does not match the
write-side range") {
+ val handler = newHandler(integrityCheckEnabled = true)
+ val (crc32PerSubPartition, bytesPerSubPartition) =
writeSideMetadata(payloads)
+ handler.recordMapPartitionCommitMetadata(
+ shuffleId,
+ mapPartitionId,
+ numPartitions,
+ crc32PerSubPartition,
+ bytesPerSubPartition)
+ val startSubIndex = 0
+ val endSubIndex = 1
+ // Flip a byte in the consumed payload to simulate a corrupted read.
+ val corrupted = payloads(endSubIndex).clone()
+ corrupted(0) = (corrupted(0) ^ 0x01).toByte
+ val actual =
+ readSideMetadata(payloads.updated(endSubIndex, corrupted),
startSubIndex, endSubIndex)
+ val (isValid, message) =
+ handler.finishPartition(shuffleId, mapPartitionId, startSubIndex,
endSubIndex, actual)
+ assert(!isValid)
+ assert(message.contains("Integrity check failed"))
+ }
+
+ test(
+ "finishPartition fails on a byte-count mismatch even when the checksum
coincidentally matches") {
+ val handler = newHandler(integrityCheckEnabled = true)
+ val (crc32PerSubPartition, bytesPerSubPartition) =
writeSideMetadata(payloads)
+ handler.recordMapPartitionCommitMetadata(
+ shuffleId,
+ mapPartitionId,
+ numPartitions,
+ crc32PerSubPartition,
+ bytesPerSubPartition)
+ val startSubIndex = 0
+ val endSubIndex = 1
+ // Expected checksum but a wrong byte count (as if a CRC32 collision hid a
short/over read).
+ val expected = new CommitMetadata()
+ (startSubIndex to endSubIndex).foreach { i =>
+ expected.addCommitData(crc32PerSubPartition(i), bytesPerSubPartition(i))
+ }
+ val actual = new CommitMetadata(expected.getChecksum, expected.getBytes +
1)
+ assert(expected.getChecksum == actual.getChecksum)
+ assert(expected.getBytes != actual.getBytes)
+ val (isValid, message) =
+ handler.finishPartition(shuffleId, mapPartitionId, startSubIndex,
endSubIndex, actual)
+ assert(!isValid)
+ assert(message.contains("Integrity check failed"))
+ }
+
+ test("recordMapPartitionCommitMetadata is safe under concurrent mappers") {
+ val handler = newHandler(integrityCheckEnabled = true)
+ val numMapPartitions = 16
+ handler.registerShuffle(
+ shuffleId,
+ numMapPartitions,
+ isSegmentGranularityVisible = true,
+ numPartitions)
+ val (crc32PerSubPartition, bytesPerSubPartition) =
writeSideMetadata(payloads)
+
+ val startLatch = new CountDownLatch(1)
+ val doneLatch = new CountDownLatch(numMapPartitions)
+ val errors = new ConcurrentLinkedQueue[Throwable]()
+ (0 until numMapPartitions).foreach { mapPartition =>
+ val thread = new Thread(new Runnable {
+ override def run(): Unit = {
+ try {
+ startLatch.await()
+ handler.recordMapPartitionCommitMetadata(
+ shuffleId,
+ mapPartition,
+ numPartitions,
+ crc32PerSubPartition,
+ bytesPerSubPartition)
+ } catch {
+ case t: Throwable => errors.add(t)
+ } finally {
+ doneLatch.countDown()
+ }
+ }
+ })
+ thread.setDaemon(true)
+ thread.start()
+ }
+ // Release all threads at once to maximize contention on the shuffle's
metadata map.
+ startLatch.countDown()
+ assert(doneLatch.await(30, TimeUnit.SECONDS), "concurrent record timed
out")
+ assert(errors.isEmpty, s"concurrent record threw:
${errors.toArray.mkString(", ")}")
+
+ // Every mapper's metadata must have been recorded and independently
validate.
+ (0 until numMapPartitions).foreach { mapPartition =>
+ val actual = readSideMetadata(payloads, 0, numPartitions - 1)
+ val (isValid, message) =
+ handler.finishPartition(shuffleId, mapPartition, 0, numPartitions - 1,
actual)
+ assert(isValid, s"map partition $mapPartition failed validation:
$message")
+ }
+ }
+}
diff --git a/common/src/main/java/org/apache/celeborn/common/CelebornCRC32.java
b/common/src/main/java/org/apache/celeborn/common/CelebornCRC32.java
index c4c71efa6..87997e3a0 100644
--- a/common/src/main/java/org/apache/celeborn/common/CelebornCRC32.java
+++ b/common/src/main/java/org/apache/celeborn/common/CelebornCRC32.java
@@ -17,6 +17,7 @@
package org.apache.celeborn.common;
+import java.nio.ByteBuffer;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.zip.CRC32;
@@ -45,6 +46,19 @@ public class CelebornCRC32 {
return (int) hashFunction.getValue();
}
+ static int compute(ByteBuffer buffer) {
+ CRC32 hashFunction = new CRC32();
+ hashFunction.update(buffer);
+ return (int) hashFunction.getValue();
+ }
+
+ static int compute(ByteBuffer first, ByteBuffer second) {
+ CRC32 hashFunction = new CRC32();
+ hashFunction.update(first);
+ hashFunction.update(second);
+ return (int) hashFunction.getValue();
+ }
+
static int combine(int first, int second) {
first =
(((byte) second + (byte) first) & 0xFF)
@@ -68,6 +82,14 @@ public class CelebornCRC32 {
addChecksum(compute(bytes, offset, length));
}
+ void addData(ByteBuffer buffer) {
+ addChecksum(compute(buffer));
+ }
+
+ void addData(ByteBuffer first, ByteBuffer second) {
+ addChecksum(compute(first, second));
+ }
+
int get() {
return current.get();
}
diff --git
a/common/src/main/java/org/apache/celeborn/common/CommitMetadata.java
b/common/src/main/java/org/apache/celeborn/common/CommitMetadata.java
index 39b84883c..814d944ba 100644
--- a/common/src/main/java/org/apache/celeborn/common/CommitMetadata.java
+++ b/common/src/main/java/org/apache/celeborn/common/CommitMetadata.java
@@ -17,6 +17,7 @@
package org.apache.celeborn.common;
+import java.nio.ByteBuffer;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicLong;
@@ -40,6 +41,23 @@ public class CommitMetadata {
this.crc.addData(rawDataBuf, offset, length);
}
+ /**
+ * Accumulates checksum and byte count over the remaining bytes of {@code
buffer} (a transient
+ * {@code nioBuffer} view, so no copy), advancing its position to the limit.
+ */
+ public void addData(ByteBuffer buffer) {
+ this.bytes.addAndGet(buffer.remaining());
+ this.crc.addData(buffer);
+ }
+
+ /**
+ * Like {@link #addData(ByteBuffer)} but over two concatenated buffers (e.g.
a split header/data).
+ */
+ public void addData(ByteBuffer first, ByteBuffer second) {
+ this.bytes.addAndGet((long) first.remaining() + second.remaining());
+ this.crc.addData(first, second);
+ }
+
public void addCommitData(CommitMetadata commitMetadata) {
addCommitData(commitMetadata.getChecksum(), commitMetadata.getBytes());
}
diff --git
a/common/src/main/java/org/apache/celeborn/common/write/PushState.java
b/common/src/main/java/org/apache/celeborn/common/write/PushState.java
index 46714c4e8..e213bf6a1 100644
--- a/common/src/main/java/org/apache/celeborn/common/write/PushState.java
+++ b/common/src/main/java/org/apache/celeborn/common/write/PushState.java
@@ -18,6 +18,7 @@
package org.apache.celeborn.common.write;
import java.io.IOException;
+import java.nio.ByteBuffer;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicReference;
@@ -136,4 +137,10 @@ public class PushState {
commitMetadataMap.computeIfAbsent(partitionId, id -> new
CommitMetadata());
commitMetadata.addDataWithOffsetAndLength(data, offset, length);
}
+
+ public void addData(int partitionId, ByteBuffer data) {
+ CommitMetadata commitMetadata =
+ commitMetadataMap.computeIfAbsent(partitionId, id -> new
CommitMetadata());
+ commitMetadata.addData(data);
+ }
}
diff --git
a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
index 32c4eb73d..f81dacc56 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -5712,7 +5712,7 @@ object CelebornConf extends Logging {
buildConf("celeborn.client.shuffle.integrityCheck.enabled")
.categories("client")
.version("0.6.1")
- .doc("When `true`, enables end-to-end integrity checks for Spark
workloads.")
+ .doc("When `true`, enables end-to-end integrity checks for Spark and
Flink workloads.")
.booleanConf
.createWithDefault(false)
diff --git
a/common/src/test/java/org/apache/celeborn/common/CelebornCRC32Test.java
b/common/src/test/java/org/apache/celeborn/common/CelebornCRC32Test.java
index 296e37d9c..c5bcc2b0a 100644
--- a/common/src/test/java/org/apache/celeborn/common/CelebornCRC32Test.java
+++ b/common/src/test/java/org/apache/celeborn/common/CelebornCRC32Test.java
@@ -20,6 +20,8 @@ package org.apache.celeborn.common;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotEquals;
+import java.nio.ByteBuffer;
+
import org.junit.Test;
// Test data is generated from https://crccalc.com/ using "CRC-32/ISO-HDLC".
@@ -63,6 +65,48 @@ public class CelebornCRC32Test {
assertEquals(2918445923L, crc.get() & 0xFFFFFFFFL);
}
+ @Test
+ public void testComputeByteBuffer() {
+ byte[] data = "testdata".getBytes();
+ // A ByteBuffer view must hash to the same value as the byte[] overload.
+ assertEquals(CelebornCRC32.compute(data),
CelebornCRC32.compute(ByteBuffer.wrap(data)));
+ // A sliced view (offset 4, length 4) must match the offset/length
overload.
+ assertEquals(
+ CelebornCRC32.compute(data, 4, 4),
CelebornCRC32.compute(ByteBuffer.wrap(data, 4, 4)));
+ }
+
+ @Test
+ public void testComputeTwoByteBuffers() {
+ byte[] first = "test".getBytes();
+ byte[] second = "data".getBytes();
+ byte[] concatenated = "testdata".getBytes();
+ // Hashing two buffers must equal hashing their concatenation.
+ assertEquals(
+ CelebornCRC32.compute(concatenated),
+ CelebornCRC32.compute(ByteBuffer.wrap(first),
ByteBuffer.wrap(second)));
+ }
+
+ @Test
+ public void testAddDataByteBuffer() {
+ byte[] data = "testdata".getBytes();
+ CelebornCRC32 fromBuffer = new CelebornCRC32();
+ fromBuffer.addData(ByteBuffer.wrap(data, 4, 4));
+ CelebornCRC32 fromBytes = new CelebornCRC32();
+ fromBytes.addData(data, 4, 4);
+ assertEquals(fromBytes.get(), fromBuffer.get());
+ }
+
+ @Test
+ public void testAddDataTwoByteBuffers() {
+ byte[] first = "test".getBytes();
+ byte[] second = "data".getBytes();
+ CelebornCRC32 fromTwoBuffers = new CelebornCRC32();
+ fromTwoBuffers.addData(ByteBuffer.wrap(first), ByteBuffer.wrap(second));
+ CelebornCRC32 fromConcatenation = new CelebornCRC32();
+ fromConcatenation.addData("testdata".getBytes(), 0, 8);
+ assertEquals(fromConcatenation.get(), fromTwoBuffers.get());
+ }
+
@Test
public void testToString() {
CelebornCRC32 crc = new CelebornCRC32(123456789);
diff --git
a/common/src/test/java/org/apache/celeborn/common/CommitMetadataTest.java
b/common/src/test/java/org/apache/celeborn/common/CommitMetadataTest.java
index 3451d386f..a228a796c 100644
--- a/common/src/test/java/org/apache/celeborn/common/CommitMetadataTest.java
+++ b/common/src/test/java/org/apache/celeborn/common/CommitMetadataTest.java
@@ -18,6 +18,10 @@
package org.apache.celeborn.common;
import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import java.nio.ByteBuffer;
import org.junit.Assert;
import org.junit.Test;
@@ -52,6 +56,89 @@ public class CommitMetadataTest {
metadata1.getChecksum());
}
+ @Test
+ public void testAddDataByteBuffer() {
+ byte[] data = "testdata".getBytes();
+ // The ByteBuffer overload (read path) must match the byte[] overload
(write path).
+ CommitMetadata fromBuffer = new CommitMetadata();
+ fromBuffer.addData(ByteBuffer.wrap(data));
+ CommitMetadata fromBytes = new CommitMetadata();
+ fromBytes.addDataWithOffsetAndLength(data, 0, data.length);
+ assertEquals(fromBytes.getBytes(), fromBuffer.getBytes());
+ assertEquals(fromBytes.getChecksum(), fromBuffer.getChecksum());
+ }
+
+ @Test
+ public void testAddDataTwoByteBuffers() {
+ byte[] header = "test".getBytes();
+ byte[] body = "data".getBytes();
+ // A split header/data buffer must hash equal to its concatenation.
+ CommitMetadata fromTwoBuffers = new CommitMetadata();
+ fromTwoBuffers.addData(ByteBuffer.wrap(header), ByteBuffer.wrap(body));
+ CommitMetadata fromConcatenation = new CommitMetadata();
+ fromConcatenation.addData(ByteBuffer.wrap("testdata".getBytes()));
+ assertEquals(fromConcatenation.getBytes(), fromTwoBuffers.getBytes());
+ assertEquals(fromConcatenation.getChecksum(),
fromTwoBuffers.getChecksum());
+ }
+
+ @Test
+ public void testRangeCombineDetectsCorruption() {
+ // Mirror the map-partition flow: per-subpartition write checksums
combined vs read bytes
+ // served.
+ byte[][] perSubpartition = {
+ "sub-0-bytes".getBytes(), "sub-1-bytes".getBytes(),
"sub-2-bytes".getBytes()
+ };
+ int startSubIndex = 0;
+ int endSubIndex = 1;
+
+ CommitMetadata expected = new CommitMetadata();
+ for (int i = startSubIndex; i <= endSubIndex; i++) {
+ CommitMetadata writeSide = new CommitMetadata();
+ writeSide.addDataWithOffsetAndLength(perSubpartition[i], 0,
perSubpartition[i].length);
+ expected.addCommitData(writeSide.getChecksum(), writeSide.getBytes());
+ }
+
+ // A faithful read over the same range matches.
+ CommitMetadata cleanRead = new CommitMetadata();
+ for (int i = startSubIndex; i <= endSubIndex; i++) {
+ cleanRead.addData(ByteBuffer.wrap(perSubpartition[i]));
+ }
+ assertTrue(CommitMetadata.checkCommitMetadata(expected, cleanRead));
+
+ // A single flipped byte on the read path is detected.
+ byte[] corrupted = perSubpartition[endSubIndex].clone();
+ corrupted[0] ^= 0x01;
+ CommitMetadata corruptedRead = new CommitMetadata();
+ corruptedRead.addData(ByteBuffer.wrap(perSubpartition[startSubIndex]));
+ corruptedRead.addData(ByteBuffer.wrap(corrupted));
+ assertFalse(CommitMetadata.checkCommitMetadata(expected, corruptedRead));
+ }
+
+ @Test
+ public void testRangeCombineIsOrderIndependent() {
+ // Tiered reads accumulate interleaved buffers while the driver combines
per-subpartition
+ // checksums in order; the combine must be order-independent for the two
to match.
+ byte[][] perSubpartition = {
+ "sub-0-bytes".getBytes(), "sub-1-bytes".getBytes(),
"sub-2-bytes".getBytes()
+ };
+
+ CommitMetadata expected = new CommitMetadata();
+ for (byte[] payload : perSubpartition) {
+ CommitMetadata writeSide = new CommitMetadata();
+ writeSide.addDataWithOffsetAndLength(payload, 0, payload.length);
+ expected.addCommitData(writeSide.getChecksum(), writeSide.getBytes());
+ }
+
+ // Read side accumulates the same bytes in reverse order; the result must
be identical.
+ CommitMetadata reversedRead = new CommitMetadata();
+ for (int i = perSubpartition.length - 1; i >= 0; i--) {
+ reversedRead.addData(ByteBuffer.wrap(perSubpartition[i]));
+ }
+ assertEquals(expected.getChecksum(), reversedRead.getChecksum());
+ assertEquals(expected.getBytes(), reversedRead.getBytes());
+ assertTrue(CommitMetadata.checkCommitMetadata(expected, reversedRead));
+ }
+
@Test
public void testCheckCommitMetadata() {
CommitMetadata expected = new
CommitMetadata(CelebornCRC32.compute("testdata".getBytes()), 8);
diff --git a/docs/configuration/client.md b/docs/configuration/client.md
index 52d3d984e..aec9fcc5f 100644
--- a/docs/configuration/client.md
+++ b/docs/configuration/client.md
@@ -114,7 +114,7 @@ license: |
| celeborn.client.shuffle.dynamicResourceEnabled | false | false | When
enabled, the ChangePartitionManager will obtain candidate workers from the
availableWorkers pool during heartbeats when worker resource change. | 0.6.0 |
|
| celeborn.client.shuffle.dynamicResourceFactor | 0.5 | false | The
ChangePartitionManager will check whether (unavailable workers / shuffle
allocated workers) is more than the factor before obtaining candidate workers
from the requestSlots RPC response when
`celeborn.client.shuffle.dynamicResourceEnabled` set true | 0.6.0 | |
| celeborn.client.shuffle.expired.checkInterval | 60s | false | Interval for
client to check expired shuffles. | 0.3.0 |
celeborn.shuffle.expired.checkInterval |
-| celeborn.client.shuffle.integrityCheck.enabled | false | false | When
`true`, enables end-to-end integrity checks for Spark workloads. | 0.6.1 | |
+| celeborn.client.shuffle.integrityCheck.enabled | false | false | When
`true`, enables end-to-end integrity checks for Spark and Flink workloads. |
0.6.1 | |
| celeborn.client.shuffle.manager.port | 0 | false | Port used by the
LifecycleManager on the Driver. | 0.3.0 | celeborn.shuffle.manager.port |
| celeborn.client.shuffle.partition.type | REDUCE | false | Type of shuffle's
partition. | 0.3.0 | celeborn.shuffle.partition.type |
| celeborn.client.shuffle.partitionSplit.mode | SOFT | false | soft: the
shuffle file size might be larger than split threshold. hard: the shuffle file
size will be limited to split threshold. | 0.3.0 |
celeborn.shuffle.partitionSplit.mode |
diff --git
a/tests/flink-it/src/test/scala/org/apache/celeborn/tests/flink/HybridShuffleWordCountTest.scala
b/tests/flink-it/src/test/scala/org/apache/celeborn/tests/flink/HybridShuffleWordCountTest.scala
index 1bfa24733..6794f23df 100644
---
a/tests/flink-it/src/test/scala/org/apache/celeborn/tests/flink/HybridShuffleWordCountTest.scala
+++
b/tests/flink-it/src/test/scala/org/apache/celeborn/tests/flink/HybridShuffleWordCountTest.scala
@@ -29,6 +29,7 @@ import
org.apache.flink.streaming.api.graph.StreamingJobGraphGenerator
import org.scalatest.BeforeAndAfterAll
import org.scalatest.funsuite.AnyFunSuite
+import org.apache.celeborn.common.CelebornConf
import org.apache.celeborn.common.internal.Logging
import org.apache.celeborn.service.deploy.MiniClusterFeature
import org.apache.celeborn.service.deploy.worker.Worker
@@ -62,12 +63,23 @@ class HybridShuffleWordCountTest extends AnyFunSuite with
Logging with MiniClust
testLocalEnv()
}
+ test("Celeborn Flink Hybrid Shuffle Integration test(Local) with integrity
check - word count") {
+ assumeFlinkVersion()
+ testLocalEnv(integrityCheckEnabled = true)
+ }
+
test(
"Celeborn Flink Hybrid Shuffle Integration test(Flink mini cluster) single
tier - word count") {
assumeFlinkVersion()
testInMiniCluster()
}
+ test(
+ "Celeborn Flink Hybrid Shuffle Integration test(Flink mini cluster) single
tier with integrity check - word count") {
+ assumeFlinkVersion()
+ testInMiniCluster(integrityCheckEnabled = true)
+ }
+
private def assumeFlinkVersion(): Unit = {
// Celeborn supports flink hybrid shuffle staring from flink 1.20
val flinkVersion = sys.env.getOrElse("FLINK_VERSION", "")
@@ -77,7 +89,7 @@ class HybridShuffleWordCountTest extends AnyFunSuite with
Logging with MiniClust
FlinkVersion.v1_20))
}
- private def testLocalEnv(): Unit = {
+ private def testLocalEnv(integrityCheckEnabled: Boolean = false): Unit = {
// set up execution environment
val configuration = new Configuration
val parallelism = NUM_PARALLELISM
@@ -88,6 +100,9 @@ class HybridShuffleWordCountTest extends AnyFunSuite with
Logging with MiniClust
"taskmanager.network.hybrid-shuffle.external-remote-tier-factory.class",
"org.apache.celeborn.plugin.flink.tiered.CelebornTierFactory")
configuration.setString("celeborn.master.endpoints", "localhost:9097")
+ if (integrityCheckEnabled) {
+
configuration.setString(CelebornConf.CLIENT_SHUFFLE_INTEGRITY_CHECK_ENABLED.key,
"true")
+ }
configuration.set(ExecutionOptions.RUNTIME_MODE,
RuntimeExecutionMode.BATCH)
configuration.setString(
"execution.batch-shuffle-mode",
@@ -115,7 +130,7 @@ class HybridShuffleWordCountTest extends AnyFunSuite with
Logging with MiniClust
checkFlushingFileLength()
}
- private def testInMiniCluster(): Unit = {
+ private def testInMiniCluster(integrityCheckEnabled: Boolean = false): Unit
= {
// set up execution environment
val configuration = new Configuration
val parallelism = NUM_PARALLELISM
@@ -126,6 +141,9 @@ class HybridShuffleWordCountTest extends AnyFunSuite with
Logging with MiniClust
"taskmanager.network.hybrid-shuffle.external-remote-tier-factory.class",
"org.apache.celeborn.plugin.flink.tiered.CelebornTierFactory")
configuration.setString("celeborn.master.endpoints", "localhost:9097")
+ if (integrityCheckEnabled) {
+
configuration.setString(CelebornConf.CLIENT_SHUFFLE_INTEGRITY_CHECK_ENABLED.key,
"true")
+ }
configuration.set(ExecutionOptions.RUNTIME_MODE,
RuntimeExecutionMode.BATCH)
configuration.setString(
"execution.batch-shuffle-mode",
diff --git
a/tests/flink-it/src/test/scala/org/apache/celeborn/tests/flink/WordCountTest.scala
b/tests/flink-it/src/test/scala/org/apache/celeborn/tests/flink/WordCountTest.scala
index fa8860efc..e3cf46191 100644
---
a/tests/flink-it/src/test/scala/org/apache/celeborn/tests/flink/WordCountTest.scala
+++
b/tests/flink-it/src/test/scala/org/apache/celeborn/tests/flink/WordCountTest.scala
@@ -127,6 +127,13 @@ class WordCountTest extends WordCountTestBase {
override protected def getWorkerConf: Map[String, String] = Map()
}
+class WordCountTestWithIntegrityCheck extends WordCountTestBase {
+ override protected def getMasterConf: Map[String, String] = Map()
+ override protected def getWorkerConf: Map[String, String] = Map()
+ override protected def getClientConf: Map[String, String] =
+ Map(CelebornConf.CLIENT_SHUFFLE_INTEGRITY_CHECK_ENABLED.key -> "true")
+}
+
class WordCountTestWithAuthentication extends WordCountTestBase {
private val authConfig = Map(
