[WIP][SPARK-35275][CORE] Add checksum for shuffle blocks and diagnose corruption

common/network-common/src/main/java/org/apache/spark/network/corruption/Cause.java

@@ -0,0 +1,22 @@
+/*
+ * 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.spark.network.corruption;
+
+public enum Cause {
+    DISK, NETWORK, UNKNOWN;
+}

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/BlockStoreClient.java

@@ -33,6 +33,7 @@
 import org.apache.spark.network.client.RpcResponseCallback;
 import org.apache.spark.network.client.TransportClient;
 import org.apache.spark.network.client.TransportClientFactory;
+import org.apache.spark.network.corruption.Cause;
 import org.apache.spark.network.shuffle.protocol.BlockTransferMessage;
 import org.apache.spark.network.shuffle.protocol.GetLocalDirsForExecutors;
 import org.apache.spark.network.shuffle.protocol.LocalDirsForExecutors;
@@ -47,6 +48,15 @@ public abstract class BlockStoreClient implements Closeable {
   protected volatile TransportClientFactory clientFactory;
   protected String appId;
 
+  public Cause diagnoseCorruption(
+      String host,
+      int port,
+      String execId,
+      String blockId,
+      long checksum) {
+    return Cause.UNKNOWN;
+  }
+
   /**
    * Fetch a sequence of blocks from a remote node asynchronously,
    *

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/ExternalBlockStoreClient.java

@@ -34,6 +34,7 @@
 import org.apache.spark.network.client.RpcResponseCallback;
 import org.apache.spark.network.client.TransportClient;
 import org.apache.spark.network.client.TransportClientBootstrap;
+import org.apache.spark.network.corruption.Cause;
 import org.apache.spark.network.crypto.AuthClientBootstrap;
 import org.apache.spark.network.sasl.SecretKeyHolder;
 import org.apache.spark.network.server.NoOpRpcHandler;
@@ -82,6 +83,16 @@ public void init(String appId) {
     clientFactory = context.createClientFactory(bootstraps);
   }
 
+  @Override
+  public Cause diagnoseCorruption(
+      String host,
+      int port,
+      String execId,
+      String blockId,
+      long checksum) {
+    return Cause.UNKNOWN;
+  }
+
   @Override
   public void fetchBlocks(
       String host,

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/protocol/BlockTransferMessage.java

@@ -48,7 +48,8 @@ public enum Type {
     OPEN_BLOCKS(0), UPLOAD_BLOCK(1), REGISTER_EXECUTOR(2), STREAM_HANDLE(3), REGISTER_DRIVER(4),
     HEARTBEAT(5), UPLOAD_BLOCK_STREAM(6), REMOVE_BLOCKS(7), BLOCKS_REMOVED(8),
     FETCH_SHUFFLE_BLOCKS(9), GET_LOCAL_DIRS_FOR_EXECUTORS(10), LOCAL_DIRS_FOR_EXECUTORS(11),
-    PUSH_BLOCK_STREAM(12), FINALIZE_SHUFFLE_MERGE(13), MERGE_STATUSES(14);
+    PUSH_BLOCK_STREAM(12), FINALIZE_SHUFFLE_MERGE(13), MERGE_STATUSES(14),
+    DIAGNOSE_CORRUPTION(15), CORRUPTION_CAUSE(16);
 
     private final byte id;
 
@@ -82,6 +83,8 @@ public static BlockTransferMessage fromByteBuffer(ByteBuffer msg) {
         case 12: return PushBlockStream.decode(buf);
         case 13: return FinalizeShuffleMerge.decode(buf);
         case 14: return MergeStatuses.decode(buf);
+        case 15: return DiagnoseCorruption.decode(buf);
+        case 16: return CorruptionCause.decode(buf);
         default: throw new IllegalArgumentException("Unknown message type: " + type);
       }
     }

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/protocol/CorruptionCause.java

@@ -0,0 +1,74 @@
+/*
+ * 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.spark.network.shuffle.protocol;
+
+import io.netty.buffer.ByteBuf;
+import org.apache.commons.lang3.builder.ToStringBuilder;
+import org.apache.commons.lang3.builder.ToStringStyle;
+import org.apache.spark.network.corruption.Cause;
+
+/** Response to the {@link DiagnoseCorruption} */
+public class CorruptionCause extends BlockTransferMessage {
+    public Cause cause;
+
+    public CorruptionCause(Cause cause) {
+      this.cause = cause;
+    }
+
+    @Override
+    protected Type type() {
+      return Type.CORRUPTION_CAUSE;
+    }
+
+    @Override
+    public String toString() {
+      return new ToStringBuilder(this, ToStringStyle.SHORT_PREFIX_STYLE)
+        .append("cause", cause)
+        .toString();
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) return true;
+      if (o == null || getClass() != o.getClass()) return false;
+
+      CorruptionCause that = (CorruptionCause) o;
+      return cause == that.cause;
+    }
+
+    @Override
+    public int hashCode() {
+      return cause.hashCode();
+    }
+
+    @Override
+    public int encodedLength() {
+      return 4; /* encoded length of cause */
+    }
+
+    @Override
+    public void encode(ByteBuf buf) {
+      buf.writeInt(cause.ordinal());
+    }
+
+    public static CorruptionCause decode(ByteBuf buf) {
+      int ordinal = buf.readInt();
+      return new CorruptionCause(Cause.values()[ordinal]);
+    }
+
+}

common/network-shuffle/src/main/java/org/apache/spark/network/shuffle/protocol/DiagnoseCorruption.java

@@ -0,0 +1,99 @@
+/*
+ * 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.spark.network.shuffle.protocol;
+
+import io.netty.buffer.ByteBuf;
+import org.apache.commons.lang3.builder.ToStringBuilder;
+import org.apache.commons.lang3.builder.ToStringStyle;
+import org.apache.spark.network.protocol.Encoders;
+
+/** Request to get the cause of a corrupted block. Returns {@link CorruptionCause} */
+public class DiagnoseCorruption extends BlockTransferMessage {
+    private final String appId;
+    private final String execId;
+    public final String blockId;
+    public final long checksum;
+
+    public DiagnoseCorruption(String appId, String execId, String blockId, long checksum) {
+      this.appId = appId;
+      this.execId = execId;
+      this.blockId = blockId;
+      this.checksum = checksum;
+    }
+
+    @Override
+    protected Type type() {
+      return Type.DIAGNOSE_CORRUPTION;
+    }
+
+    @Override
+    public String toString() {
+      return new ToStringBuilder(this, ToStringStyle.SHORT_PREFIX_STYLE)
+        .append("appId", appId)
+        .append("execId", execId)
+        .append("blockId", blockId)
+        .append("checksum", checksum)
+        .toString();
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) return true;
+      if (o == null || getClass() != o.getClass()) return false;
+
+      DiagnoseCorruption that = (DiagnoseCorruption) o;
+
+      if (!appId.equals(that.appId)) return false;
+      if (!execId.equals(that.execId)) return false;
+      if (!blockId.equals(that.blockId)) return false;
+      return checksum == that.checksum;
+    }
+
+    @Override
+    public int hashCode() {
+      int result = appId.hashCode();
+      result = 31 * result + execId.hashCode();
+      result = 31 * result + blockId.hashCode();
+      result = 31 * result + (int) checksum;
+      return result;
+    }
+
+    @Override
+    public int encodedLength() {
+      return Encoders.Strings.encodedLength(appId)
+        + Encoders.Strings.encodedLength(execId)
+        + Encoders.Strings.encodedLength(blockId)
+        + 8; /* encoded length of checksum */
+    }
+
+    @Override
+    public void encode(ByteBuf buf) {
+      Encoders.Strings.encode(buf, appId);
+      Encoders.Strings.encode(buf, execId);
+      Encoders.Strings.encode(buf, blockId);
+      buf.writeLong(checksum);
+    }
+
+    public static DiagnoseCorruption decode(ByteBuf buf) {
+      String appId = Encoders.Strings.decode(buf);
+      String execId = Encoders.Strings.decode(buf);
+      String blockId = Encoders.Strings.decode(buf);
+      long checksum = buf.readLong();
+      return new DiagnoseCorruption(appId, execId, blockId, checksum);
+    }
+}

core/src/main/java/org/apache/spark/io/CountingWritableChannel.java

@@ -0,0 +1,65 @@
+/*
+ * 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.spark.io;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.WritableByteChannel;
+
+import org.apache.spark.annotation.Private;
+
+/**
+ * :: Private ::
+ * A wrapper for {@link java.nio.channels.WritableByteChannel} with the ability of counting
+ * written bytes.
+ */
+@Private
+public class CountingWritableChannel implements WritableByteChannel {
+
+  private WritableByteChannel delegate;
+
+  private long count;
+
+  public CountingWritableChannel(WritableByteChannel delegate) {
+    this.delegate = delegate;
+    this.count = 0;
+  }
+
+  public long getCount() {
+    return this.count;
+  }
+
+  @Override
+  public int write(ByteBuffer src) throws IOException {
+    int written = delegate.write(src);
+    if (written > 0) {
+      count += written;
+    }
+    return written;
+  }
+
+  @Override
+  public boolean isOpen() {
+    return delegate.isOpen();
+  }
+
+  @Override
+  public void close() throws IOException {
+    delegate.close();
+  }
+}

core/src/main/java/org/apache/spark/shuffle/api/SingleSpillShuffleMapOutputWriter.java

@@ -32,5 +32,8 @@ public interface SingleSpillShuffleMapOutputWriter {
   /**
    * Transfer a file that contains the bytes of all the partitions written by this map task.
    */
-  void transferMapSpillFile(File mapOutputFile, long[] partitionLengths) throws IOException;
+  void transferMapSpillFile(
+      File mapOutputFile,
+      long[] partitionLengths,
+      long[] checksums) throws IOException;
 }

core/src/main/java/org/apache/spark/shuffle/api/metadata/MapOutputCommitMessage.java

@@ -36,27 +36,39 @@
 public final class MapOutputCommitMessage {
 
   private final long[] partitionLengths;
+  private final long[] partitionChecksums;
   private final Optional<MapOutputMetadata> mapOutputMetadata;
 
   private MapOutputCommitMessage(
-      long[] partitionLengths, Optional<MapOutputMetadata> mapOutputMetadata) {
+      long[] partitionLengths,
+      long[] partitionChecksums,
+      Optional<MapOutputMetadata> mapOutputMetadata) {
     this.partitionLengths = partitionLengths;
+    this.partitionChecksums = partitionChecksums;
     this.mapOutputMetadata = mapOutputMetadata;
   }
 
   public static MapOutputCommitMessage of(long[] partitionLengths) {
-    return new MapOutputCommitMessage(partitionLengths, Optional.empty());
+    return new MapOutputCommitMessage(partitionLengths, null, Optional.empty());
+  }
+
+  public static MapOutputCommitMessage of(long[] partitionLengths, long[] partitionChecksums) {
+    return new MapOutputCommitMessage(partitionLengths, partitionChecksums, Optional.empty());
   }
 
   public static MapOutputCommitMessage of(
       long[] partitionLengths, MapOutputMetadata mapOutputMetadata) {
-    return new MapOutputCommitMessage(partitionLengths, Optional.of(mapOutputMetadata));
+    return new MapOutputCommitMessage(partitionLengths, null, Optional.of(mapOutputMetadata));
   }
 
   public long[] getPartitionLengths() {
     return partitionLengths;
   }
 
+  public long[] getPartitionChecksums() {
+    return partitionChecksums;
+  }
+
   public Optional<MapOutputMetadata> getMapOutputMetadata() {
     return mapOutputMetadata;
   }

core/src/main/java/org/apache/spark/shuffle/sort/BypassMergeSortShuffleWriter.java

@@ -21,6 +21,7 @@
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.OutputStream;
+import java.nio.ByteBuffer;
 import java.nio.channels.FileChannel;
 import java.util.Optional;
 import javax.annotation.Nullable;