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 e5099f90d [CELEBORN-2313] Extend E2E checked zone to batch assembly
point
e5099f90d is described below
commit e5099f90d55e7b70292b0677c3a889f72c9b7efb
Author: James Xu <[email protected]>
AuthorDate: Tue Jun 9 03:18:57 2026 +0800
[CELEBORN-2313] Extend E2E checked zone to batch assembly point
### What changes were proposed in this pull request?
Celeborn's E2E integrity check computes CRC_M inside
`ShuffleClientImpl.pushOrMergeData()`, which runs in the async `DataPusher`
thread. This leaves the segment from batch assembly in the writer thread
through the `DataPusher` queue entirely outside the checked zone — meaning any
corruption that occurs in that window is invisible to the integrity check and
reaches reducers silently.
This change closes that gap and enables detection of a class of correctness
bugs where data corruption occurs between batch assembly and async push
dispatch, including bugs involving shared buffer pool references.
Introduce `ShuffleClient.computeBatchCRC()` and consolidate its invocation
into two choke points:
- `DataPusher.addTask()`: covers all async push paths. The CRC is recorded
on the writer thread immediately before the buffer is enqueued, so
`DataPusher.pushData()` intentionally uses the bare `client.pushData()` to
avoid double-counting the same batch into CommitMetadata.
- `ShuffleClient.pushDataWithCRC()` and `ShuffleClient.mergeDataWithCRC()`:
new concrete convenience methods that call `computeBatchCRC()` then delegate to
the abstract `pushData()`/`mergeData()`. These cover all synchronous push paths
(`pushGiantRecord`, `close()` flush). The abstract `pushData()`/`mergeData()`
are now documented as internal-use-only; all writer call sites across
`HashBasedShuffleWriter` (spark-2/3), `SortBasedShuffleWriter` (spark-2/3), and
`SortBasedPusher` use the [...]
The now-redundant CRC computation inside `pushOrMergeData()` is removed.
This consolidation eliminates 7 scattered `computeBatchCRC` call sites that
previously had to be manually paired with each push/merge call, reducing the
risk of a future call site omitting the CRC step.
### Why are the changes needed?
Enhance E2E Integrity Check, so it can cover more code path.
### Does this PR resolve a correctness bug?
- [ ] Yes
### Does this PR introduce _any_ user-facing change?
- [ ] Yes
### How was this patch tested?
Unit Test.
Closes #3716 from xumingming/extend-e2e-checked-zone-v2.
Authored-by: James Xu <[email protected]>
Signed-off-by: Nicholas Jiang <[email protected]>
---
.../spark/shuffle/celeborn/SortBasedPusher.java | 2 +-
.../shuffle/celeborn/HashBasedShuffleWriter.java | 4 +-
.../shuffle/celeborn/SortBasedShuffleWriter.java | 2 +-
.../shuffle/celeborn/HashBasedShuffleWriter.java | 4 +-
.../shuffle/celeborn/SortBasedShuffleWriter.java | 2 +-
.../celeborn/CelebornShuffleWriterSuiteBase.java | 93 ++++++++++++++++++
.../apache/celeborn/client/DummyShuffleClient.java | 59 +++++++++++-
.../org/apache/celeborn/client/ShuffleClient.java | 106 ++++++++++++++++++++-
.../apache/celeborn/client/ShuffleClientImpl.java | 26 +++--
.../apache/celeborn/client/write/DataPusher.java | 3 +
.../celeborn/client/ShuffleClientSuiteJ.java | 97 +++++++++++++++++++
11 files changed, 383 insertions(+), 15 deletions(-)
diff --git
a/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedPusher.java
b/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedPusher.java
index 013785ecc..8297aefb6 100644
---
a/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedPusher.java
+++
b/client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedPusher.java
@@ -222,7 +222,7 @@ public class SortBasedPusher extends MemoryConsumer {
currentPartition = partition;
} else {
int bytesWritten =
- shuffleClient.mergeData(
+ shuffleClient.mergeDataWithCRC(
shuffleId,
mapId,
attemptNumber,
diff --git
a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
index de7d00b4e..4d55c14be 100644
---
a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
+++
b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
@@ -278,7 +278,7 @@ public class HashBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
private void pushGiantRecord(int partitionId, byte[] buffer, int numBytes)
throws IOException {
logger.debug("Push giant record for partition {}, size {}.", partitionId,
numBytes);
int bytesWritten =
- shuffleClient.pushData(
+ shuffleClient.pushDataWithCRC(
shuffleId,
mapId,
encodedAttemptId,
@@ -339,7 +339,7 @@ public class HashBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
final int size = sendOffsets[i];
if (size > 0) {
int bytesWritten =
- shuffleClient.mergeData(
+ shuffleClient.mergeDataWithCRC(
shuffleId,
mapId,
encodedAttemptId,
diff --git
a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
index 9ba908ade..2ab3139b5 100644
---
a/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
+++
b/client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
@@ -284,7 +284,7 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
private void pushGiantRecord(int partitionId, byte[] buffer, int numBytes)
throws IOException {
logger.debug("Push giant record, size {}.", Utils.bytesToString(numBytes));
int bytesWritten =
- shuffleClient.pushData(
+ shuffleClient.pushDataWithCRC(
shuffleId,
mapId,
encodedAttemptId,
diff --git
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
index 49c6d6954..8b454ebd5 100644
---
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
+++
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java
@@ -279,7 +279,7 @@ public class HashBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
logger.debug("Push giant record, size {}.", numBytes);
long start = System.nanoTime();
int bytesWritten =
- shuffleClient.pushData(
+ shuffleClient.pushDataWithCRC(
shuffleId,
mapId,
encodedAttemptId,
@@ -343,7 +343,7 @@ public class HashBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
protected void mergeData(int partitionId, byte[] buffer, int offset, int
length)
throws IOException {
int bytesWritten =
- shuffleClient.mergeData(
+ shuffleClient.mergeDataWithCRC(
shuffleId,
mapId,
encodedAttemptId,
diff --git
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
index e413ce42f..ebc1d44ea 100644
---
a/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
+++
b/client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java
@@ -349,7 +349,7 @@ public class SortBasedShuffleWriter<K, V, C> extends
ShuffleWriter<K, V> {
logger.debug("Push giant record, size {}.",
Utils.bytesToString(numBytes));
long start = System.nanoTime();
int bytesWritten =
- shuffleClient.pushData(
+ shuffleClient.pushDataWithCRC(
shuffleId,
mapId,
encodedAttemptId,
diff --git
a/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/CelebornShuffleWriterSuiteBase.java
b/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/CelebornShuffleWriterSuiteBase.java
index 62d31ee6d..8c74e0946 100644
---
a/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/CelebornShuffleWriterSuiteBase.java
+++
b/client-spark/spark-3/src/test/java/org/apache/spark/shuffle/celeborn/CelebornShuffleWriterSuiteBase.java
@@ -26,6 +26,8 @@ import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.util.Collections;
+import java.util.List;
+import java.util.Map;
import java.util.Random;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicInteger;
@@ -78,9 +80,11 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.client.DummyShuffleClient;
import org.apache.celeborn.client.ShuffleClient;
import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.CommitMetadata;
import org.apache.celeborn.common.identity.UserIdentifier;
import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.Utils;
+import org.apache.celeborn.common.write.PushState;
import org.apache.celeborn.reflect.DynConstructors;
public abstract class CelebornShuffleWriterSuiteBase {
@@ -213,6 +217,95 @@ public abstract class CelebornShuffleWriterSuiteBase {
check(2 << 30, conf, serializer);
}
+ @Test
+ public void testIntegrityCheckAccumulation() throws Exception {
+ final KryoSerializer serializer = new KryoSerializer(sparkConf);
+ final CelebornConf conf =
+ new CelebornConf()
+ .set(CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE().key(), "128")
+ .set("celeborn.client.shuffle.integrityCheck.enabled", "true");
+ checkWithIntegrity(10000, conf, serializer);
+ }
+
+ @Test
+ public void testIntegrityCheckAccumulationWithFastWrite() throws Exception {
+ final UnsafeRowSerializer serializer = new UnsafeRowSerializer(2, null);
+ final CelebornConf conf =
+ new CelebornConf()
+ .set(CelebornConf.CLIENT_PUSH_BUFFER_MAX_SIZE().key(), "128")
+ .set("celeborn.client.shuffle.integrityCheck.enabled", "true");
+ checkWithIntegrity(10000, conf, serializer);
+ }
+
+ private void checkWithIntegrity(
+ final int approximateSize, final CelebornConf conf, final Serializer
serializer)
+ throws Exception {
+ final boolean useUnsafe = serializer instanceof UnsafeRowSerializer;
+
+ String partitionIdPassthroughClazz;
+ if (SparkVersionUtil.isGreaterThan(3, 3)) {
+ partitionIdPassthroughClazz = "org.apache.spark.PartitionIdPassthrough";
+ } else {
+ partitionIdPassthroughClazz =
"org.apache.spark.sql.execution.PartitionIdPassthrough";
+ }
+ DynConstructors.Ctor<Partitioner> partitionIdPassthroughCtor =
+ DynConstructors.builder().impl(partitionIdPassthroughClazz,
int.class).build();
+ final Partitioner partitioner =
+ useUnsafe
+ ? partitionIdPassthroughCtor.newInstance(numPartitions)
+ : new HashPartitioner(numPartitions);
+ Mockito.doReturn(partitioner).when(dependency).partitioner();
+ Mockito.doReturn(serializer).when(dependency).serializer();
+
+ final File tempFile = new File(tempDir, UUID.randomUUID().toString());
+ final DummyShuffleClient client = new DummyShuffleClient(conf, tempFile);
+ client.initReducePartitionMap(shuffleId, numPartitions, 1);
+
+ final CelebornShuffleHandle<Integer, String, String> handle =
+ new CelebornShuffleHandle<>(
+ appId, host, port, userIdentifier, shuffleId, false, numMaps,
dependency);
+ final ShuffleWriter<Integer, String> writer =
+ createShuffleWriter(handle, taskContext, conf, client,
metrics.shuffleWriteMetrics());
+
+ AtomicInteger total = new AtomicInteger(0);
+ // Use mix=true to exercise both giant and normal record paths.
+ Iterator iterator = getIterator(approximateSize, total, useUnsafe, true);
+
+ writer.write(iterator);
+ Option<MapStatus> status = writer.stop(true);
+
+ assertNotNull(status);
+ assertTrue(status.isDefined());
+
+ // mapId=0 and attemptId=0 from the mock TaskContext.
+ String mapKey = Utils.makeMapKey(shuffleId, 0, 0);
+ PushState pushState = client.getPushState(mapKey);
+
+ int[] crcPerPartition = pushState.getCRC32PerPartition(true,
numPartitions);
+ long[] bytesPerPartition = pushState.getBytesWrittenPerPartition(true,
numPartitions);
+
+ Map<Integer, List<byte[]>> crcData = client.getCrcDataByPartition(mapKey);
+ Map<Integer, List<byte[]>> pushData =
client.getPushDataByPartition(mapKey);
+
+ for (int i = 0; i < numPartitions; i++) {
+ List<byte[]> crcBatches = crcData.getOrDefault(i,
Collections.emptyList());
+ long expectedBytes = 0;
+ CommitMetadata expected = new CommitMetadata();
+ for (byte[] batch : crcBatches) {
+ expected.addDataWithOffsetAndLength(batch, 0, batch.length);
+ expectedBytes += batch.length;
+ }
+ assertEquals("Partition " + i + " bytes mismatch", expectedBytes,
bytesPerPartition[i]);
+ assertEquals("Partition " + i + " CRC mismatch", expected.getChecksum(),
crcPerPartition[i]);
+
+ int pushCalls = pushData.getOrDefault(i, Collections.emptyList()).size();
+ assertEquals(
+ "Partition " + i + " CRC/push call count mismatch", pushCalls,
crcBatches.size());
+ }
+
+ client.shutdown();
+ }
+
private void check(
final int approximateSize, final CelebornConf conf, final Serializer
serializer)
throws Exception {
diff --git
a/client/src/main/java/org/apache/celeborn/client/DummyShuffleClient.java
b/client/src/main/java/org/apache/celeborn/client/DummyShuffleClient.java
index 6ca3406be..8390a9c1b 100644
--- a/client/src/main/java/org/apache/celeborn/client/DummyShuffleClient.java
+++ b/client/src/main/java/org/apache/celeborn/client/DummyShuffleClient.java
@@ -25,6 +25,8 @@ import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
@@ -47,6 +49,7 @@ import org.apache.celeborn.common.protocol.PbStreamHandler;
import org.apache.celeborn.common.rpc.RpcEndpointRef;
import org.apache.celeborn.common.util.ExceptionMaker;
import org.apache.celeborn.common.util.JavaUtils;
+import org.apache.celeborn.common.util.Utils;
import org.apache.celeborn.common.write.LocationPushFailedBatches;
import org.apache.celeborn.common.write.PushState;
@@ -56,15 +59,24 @@ public class DummyShuffleClient extends ShuffleClient {
private final OutputStream os;
private final CelebornConf conf;
+ private final boolean shuffleIntegrityCheckEnabled;
private final Map<Integer, ConcurrentHashMap<Integer, PartitionLocation>>
reducePartitionMap =
new HashMap<>();
+ // Tracking for CRC verification in tests
+ private final Map<String, PushState> pushStateMap =
JavaUtils.newConcurrentHashMap();
+ private final Map<String, Map<Integer, List<byte[]>>> crcDataByMapKey =
+ JavaUtils.newConcurrentHashMap();
+ private final Map<String, Map<Integer, List<byte[]>>> pushDataByMapKey =
+ JavaUtils.newConcurrentHashMap();
+
public AtomicInteger fetchFailureCount = new AtomicInteger();
public DummyShuffleClient(CelebornConf conf, File file) throws Exception {
this.os = new BufferedOutputStream(new FileOutputStream(file));
this.conf = conf;
+ this.shuffleIntegrityCheckEnabled =
conf.clientShuffleIntegrityCheckEnabled();
}
@Override
@@ -88,10 +100,40 @@ public class DummyShuffleClient extends ShuffleClient {
int numMappers,
int numPartitions)
throws IOException {
+ String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
+ Map<Integer, List<byte[]>> partitionData =
+ pushDataByMapKey.computeIfAbsent(mapKey, k ->
JavaUtils.newConcurrentHashMap());
+ partitionData
+ .computeIfAbsent(partitionId, k -> Collections.synchronizedList(new
ArrayList<>()))
+ .add(Arrays.copyOfRange(data, offset, offset + length));
+
os.write(data, offset, length);
return length;
}
+ @Override
+ public void computeBatchCRC(
+ int shuffleId,
+ int mapId,
+ int attemptId,
+ int partitionId,
+ byte[] data,
+ int offset,
+ int length) {
+ if (!shuffleIntegrityCheckEnabled) {
+ return;
+ }
+ String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
+ PushState pushState = pushStateMap.computeIfAbsent(mapKey, k -> new
PushState(conf));
+ pushState.addDataWithOffsetAndLength(partitionId, data, offset, length);
+
+ Map<Integer, List<byte[]>> partitionData =
+ crcDataByMapKey.computeIfAbsent(mapKey, k ->
JavaUtils.newConcurrentHashMap());
+ partitionData
+ .computeIfAbsent(partitionId, k -> Collections.synchronizedList(new
ArrayList<>()))
+ .add(Arrays.copyOfRange(data, offset, offset + length));
+ }
+
@Override
public int mergeData(
int shuffleId,
@@ -104,6 +146,13 @@ public class DummyShuffleClient extends ShuffleClient {
int numMappers,
int numPartitions)
throws IOException {
+ String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
+ Map<Integer, List<byte[]>> partitionData =
+ pushDataByMapKey.computeIfAbsent(mapKey, k ->
JavaUtils.newConcurrentHashMap());
+ partitionData
+ .computeIfAbsent(partitionId, k -> Collections.synchronizedList(new
ArrayList<>()))
+ .add(Arrays.copyOfRange(data, offset, offset + length));
+
os.write(data, offset, length);
return length;
}
@@ -188,7 +237,15 @@ public class DummyShuffleClient extends ShuffleClient {
@Override
public PushState getPushState(String mapKey) {
- return new PushState(conf);
+ return pushStateMap.computeIfAbsent(mapKey, k -> new PushState(conf));
+ }
+
+ public Map<Integer, List<byte[]>> getCrcDataByPartition(String mapKey) {
+ return crcDataByMapKey.getOrDefault(mapKey, Collections.emptyMap());
+ }
+
+ public Map<Integer, List<byte[]>> getPushDataByPartition(String mapKey) {
+ return pushDataByMapKey.getOrDefault(mapKey, Collections.emptyMap());
}
@Override
diff --git a/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
b/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
index 7a89b051d..7035478eb 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClient.java
@@ -161,7 +161,44 @@ public abstract class ShuffleClient {
public abstract void setExtension(byte[] extension);
/**
- * Write data to a specific reduce partition
+ * Write data to a specific reduce partition, computing and recording a CRC
over the batch before
+ * pushing. Prefer this over {@link #pushData} at all writer call sites.
+ *
+ * @param shuffleId the unique shuffle id of the application
+ * @param mapId the map id of the shuffle
+ * @param attemptId the attempt id of the map task, i.e. speculative task or
task rerun for Apache
+ * Spark
+ * @param partitionId the partition id the data belongs to
+ * @param data byte array containing data to be pushed
+ * @param offset start position of data to be pushed
+ * @param length length of data to be pushed
+ * @param numMappers the number map tasks in the shuffle
+ * @param numPartitions the number of partitions in the shuffle
+ * @return bytes pushed
+ * @throws IOException
+ */
+ public int pushDataWithCRC(
+ int shuffleId,
+ int mapId,
+ int attemptId,
+ int partitionId,
+ byte[] data,
+ int offset,
+ int length,
+ int numMappers,
+ int numPartitions)
+ throws IOException {
+ computeBatchCRC(shuffleId, mapId, attemptId, partitionId, data, offset,
length);
+ return pushData(
+ shuffleId, mapId, attemptId, partitionId, data, offset, length,
numMappers, numPartitions);
+ }
+
+ /**
+ * Write data to a specific reduce partition.
+ *
+ * <p><b>Internal use only.</b> Callers outside the async push pipeline
(i.e. {@link
+ * org.apache.celeborn.client.write.DataPusher}) should use {@link
#pushDataWithCRC} instead,
+ * which additionally records a CRC over the batch for end-to-end integrity
checking.
*
* @param shuffleId the unique shuffle id of the application
* @param mapId the map id of the shuffle
@@ -188,6 +225,73 @@ public abstract class ShuffleClient {
int numPartitions)
throws IOException;
+ /**
+ * Pre-compute CRC for a batch immediately after assembly in the writer,
before the data enters
+ * the async push pipeline. This is the sole CRC accumulation path when
shuffle integrity check is
+ * enabled.
+ */
+ public abstract void computeBatchCRC(
+ int shuffleId,
+ int mapId,
+ int attemptId,
+ int partitionId,
+ byte[] data,
+ int offset,
+ int length);
+
+ /**
+ * Merge data into a specific reduce partition, computing and recording a
CRC over the batch
+ * before merging. Prefer this over {@link #mergeData} at all writer call
sites.
+ *
+ * @param shuffleId the unique shuffle id of the application
+ * @param mapId the map id of the shuffle
+ * @param attemptId the attempt id of the map task, i.e. speculative task or
task rerun for Apache
+ * Spark
+ * @param partitionId the partition id the data belongs to
+ * @param data byte array containing data to be merged
+ * @param offset start position of data to be merged
+ * @param length length of data to be merged
+ * @param numMappers the number map tasks in the shuffle
+ * @param numPartitions the number of partitions in the shuffle
+ * @return bytes merged
+ * @throws IOException
+ */
+ public int mergeDataWithCRC(
+ int shuffleId,
+ int mapId,
+ int attemptId,
+ int partitionId,
+ byte[] data,
+ int offset,
+ int length,
+ int numMappers,
+ int numPartitions)
+ throws IOException {
+ computeBatchCRC(shuffleId, mapId, attemptId, partitionId, data, offset,
length);
+ return mergeData(
+ shuffleId, mapId, attemptId, partitionId, data, offset, length,
numMappers, numPartitions);
+ }
+
+ /**
+ * Merge data into a specific reduce partition.
+ *
+ * <p><b>Internal use only.</b> Callers outside the async push pipeline
should use {@link
+ * #mergeDataWithCRC} instead, which additionally records a CRC over the
batch for end-to-end
+ * integrity checking.
+ *
+ * @param shuffleId the unique shuffle id of the application
+ * @param mapId the map id of the shuffle
+ * @param attemptId the attempt id of the map task, i.e. speculative task or
task rerun for Apache
+ * Spark
+ * @param partitionId the partition id the data belongs to
+ * @param data byte array containing data to be merged
+ * @param offset start position of data to be merged
+ * @param length length of data to be merged
+ * @param numMappers the number map tasks in the shuffle
+ * @param numPartitions the number of partitions in the shuffle
+ * @return bytes merged
+ * @throws IOException
+ */
public abstract int mergeData(
int shuffleId,
int mapId,
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 225a9cb99..358bc227a 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
@@ -1044,12 +1044,6 @@ public class ShuffleClientImpl extends ShuffleClient {
// increment batchId
final int nextBatchId = pushState.nextBatchId();
- // Track commit metadata if shuffle compression and integrity check are
enabled and this request
- // is not for pushing metadata itself.
- if (shuffleIntegrityCheckEnabled) {
- pushState.addDataWithOffsetAndLength(partitionId, data, offset, length);
- }
-
if (shuffleCompressionEnabled && !skipCompress) {
// compress data
final Compressor compressor = compressorThreadLocal.get();
@@ -1404,6 +1398,26 @@ public class ShuffleClientImpl extends ShuffleClient {
false);
}
+ @Override
+ public void computeBatchCRC(
+ int shuffleId,
+ int mapId,
+ int attemptId,
+ int partitionId,
+ byte[] data,
+ int offset,
+ int length) {
+ if (!shuffleIntegrityCheckEnabled) {
+ return;
+ }
+ final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
+ if (mapperEnded(shuffleId, mapId)) {
+ return;
+ }
+ PushState pushState = getPushState(mapKey);
+ pushState.addDataWithOffsetAndLength(partitionId, data, offset, length);
+ }
+
@Override
public void pushMergedData(int shuffleId, int mapId, int attemptId) throws
IOException {
final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
diff --git
a/client/src/main/java/org/apache/celeborn/client/write/DataPusher.java
b/client/src/main/java/org/apache/celeborn/client/write/DataPusher.java
index 73af67b07..0677502ab 100644
--- a/client/src/main/java/org/apache/celeborn/client/write/DataPusher.java
+++ b/client/src/main/java/org/apache/celeborn/client/write/DataPusher.java
@@ -153,6 +153,7 @@ public class DataPusher {
public void addTask(int partitionId, byte[] buffer, int size)
throws IOException, InterruptedException {
+ client.computeBatchCRC(shuffleId, mapId, attemptId, partitionId, buffer,
0, size);
try {
PushTask task = null;
while (task == null) {
@@ -208,6 +209,8 @@ public class DataPusher {
}
protected void pushData(PushTask task) throws IOException {
+ // CRC was already recorded in addTask() on the writer thread before the
buffer was enqueued.
+ // Use the bare pushData here to avoid double-counting the same batch into
CommitMetadata.
int bytesWritten =
client.pushData(
shuffleId,
diff --git
a/client/src/test/java/org/apache/celeborn/client/ShuffleClientSuiteJ.java
b/client/src/test/java/org/apache/celeborn/client/ShuffleClientSuiteJ.java
index e6d450d87..74fe6379c 100644
--- a/client/src/test/java/org/apache/celeborn/client/ShuffleClientSuiteJ.java
+++ b/client/src/test/java/org/apache/celeborn/client/ShuffleClientSuiteJ.java
@@ -47,6 +47,7 @@ import org.mockito.ArgumentCaptor;
import org.apache.celeborn.client.compress.Compressor;
import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.CommitMetadata;
import org.apache.celeborn.common.exception.CelebornIOException;
import org.apache.celeborn.common.identity.UserIdentifier;
import org.apache.celeborn.common.network.client.TransportClient;
@@ -61,6 +62,8 @@ import
org.apache.celeborn.common.protocol.message.ControlMessages.RegisterShuff
import org.apache.celeborn.common.protocol.message.StatusCode;
import org.apache.celeborn.common.rpc.RpcEndpointRef;
import org.apache.celeborn.common.rpc.RpcTimeoutException;
+import org.apache.celeborn.common.util.Utils;
+import org.apache.celeborn.common.write.PushState;
public class ShuffleClientSuiteJ {
@@ -709,4 +712,98 @@ public class ShuffleClientSuiteJ {
assertEquals(crc32, capturedRequest.getCrc32());
assertEquals(bytesWritten, capturedRequest.getBytesWritten());
}
+
+ @Test
+ public void testComputeBatchCRCAccumulatesCorrectly() {
+ CelebornConf conf = new CelebornConf();
+ conf.set("celeborn.client.shuffle.integrityCheck.enabled", "true");
+ shuffleClient =
+ new ShuffleClientImpl(TEST_APPLICATION_ID, conf, new
UserIdentifier("mock", "mock"));
+ shuffleClient.setupLifecycleManagerRef(endpointRef);
+
+ byte[] batch0 = "hello world".getBytes(StandardCharsets.UTF_8);
+ byte[] batch1a = "foo".getBytes(StandardCharsets.UTF_8);
+ byte[] batch1b = "bar".getBytes(StandardCharsets.UTF_8);
+
+ shuffleClient.computeBatchCRC(
+ TEST_SHUFFLE_ID, TEST_MAP_ID, TEST_ATTEMPT_ID, 0, batch0, 0,
batch0.length);
+ shuffleClient.computeBatchCRC(
+ TEST_SHUFFLE_ID, TEST_MAP_ID, TEST_ATTEMPT_ID, 1, batch1a, 0,
batch1a.length);
+ shuffleClient.computeBatchCRC(
+ TEST_SHUFFLE_ID, TEST_MAP_ID, TEST_ATTEMPT_ID, 1, batch1b, 0,
batch1b.length);
+
+ PushState pushState =
+ shuffleClient.getPushState(Utils.makeMapKey(TEST_SHUFFLE_ID,
TEST_MAP_ID, TEST_ATTEMPT_ID));
+
+ int numPartitions = 2;
+ int[] crcPerPartition = pushState.getCRC32PerPartition(true,
numPartitions);
+ long[] bytesPerPartition = pushState.getBytesWrittenPerPartition(true,
numPartitions);
+
+ // compute expected values via CommitMetadata — same code path as
production
+ CommitMetadata expected0 = new CommitMetadata();
+ expected0.addDataWithOffsetAndLength(batch0, 0, batch0.length);
+ assertEquals(expected0.getChecksum(), crcPerPartition[0]);
+ assertEquals(expected0.getBytes(), bytesPerPartition[0]);
+
+ CommitMetadata expected1 = new CommitMetadata();
+ expected1.addDataWithOffsetAndLength(batch1a, 0, batch1a.length);
+ expected1.addDataWithOffsetAndLength(batch1b, 0, batch1b.length);
+ assertEquals(expected1.getChecksum(), crcPerPartition[1]);
+ assertEquals(expected1.getBytes(), bytesPerPartition[1]);
+ }
+
+ @Test
+ public void testComputeBatchCRCDisabled() {
+ CelebornConf conf = new CelebornConf();
+ conf.set("celeborn.client.shuffle.integrityCheck.enabled", "false");
+ shuffleClient =
+ new ShuffleClientImpl(TEST_APPLICATION_ID, conf, new
UserIdentifier("mock", "mock"));
+ shuffleClient.setupLifecycleManagerRef(endpointRef);
+
+ byte[] data = "hello world".getBytes(StandardCharsets.UTF_8);
+ shuffleClient.computeBatchCRC(
+ TEST_SHUFFLE_ID, TEST_MAP_ID, TEST_ATTEMPT_ID, 0, data, 0,
data.length);
+
+ PushState pushState =
+ shuffleClient.getPushState(Utils.makeMapKey(TEST_SHUFFLE_ID,
TEST_MAP_ID, TEST_ATTEMPT_ID));
+
+ int numPartitions = 2;
+ // When integrity check is disabled at the client level, computeBatchCRC
is a no-op,
+ // so the PushState should remain empty even though we pass true to the
getter.
+ int[] crcPerPartition = pushState.getCRC32PerPartition(true,
numPartitions);
+ long[] bytesPerPartition = pushState.getBytesWrittenPerPartition(true,
numPartitions);
+
+ for (int i = 0; i < numPartitions; i++) {
+ assertEquals("Partition " + i + " CRC should be zero", 0,
crcPerPartition[i]);
+ assertEquals("Partition " + i + " bytes should be zero", 0,
bytesPerPartition[i]);
+ }
+ }
+
+ @Test
+ public void testComputeBatchCRCAttemptIdConsistency() {
+ CelebornConf conf = new CelebornConf();
+ conf.set("celeborn.client.shuffle.integrityCheck.enabled", "true");
+ shuffleClient =
+ new ShuffleClientImpl(TEST_APPLICATION_ID, conf, new
UserIdentifier("mock", "mock"));
+ shuffleClient.setupLifecycleManagerRef(endpointRef);
+
+ byte[] data = "test".getBytes(StandardCharsets.UTF_8);
+
+ // Call with attemptId=0
+ shuffleClient.computeBatchCRC(
+ TEST_SHUFFLE_ID, TEST_MAP_ID, TEST_ATTEMPT_ID, 0, data, 0,
data.length);
+
+ // PushState for attemptId=0 should have the data
+ PushState pushState0 =
+ shuffleClient.getPushState(Utils.makeMapKey(TEST_SHUFFLE_ID,
TEST_MAP_ID, TEST_ATTEMPT_ID));
+ int[] crc0 = pushState0.getCRC32PerPartition(true, 2);
+ assertTrue(crc0[0] != 0);
+
+ // PushState for attemptId=1 should be empty
+ PushState pushState1 =
+ shuffleClient.getPushState(
+ Utils.makeMapKey(TEST_SHUFFLE_ID, TEST_MAP_ID, TEST_ATTEMPT_ID +
1));
+ int[] crc1 = pushState1.getCRC32PerPartition(true, 2);
+ assertEquals(0, crc1[0]);
+ }
}
