BACKPORT: [apache#1608][part-6][FOLLOWUP] improvement(client): Check …

…blockId num after blocks all sent apache#1761

apache#1761

client-spark/spark3/src/main/java/org/apache/spark/shuffle/writer/RssShuffleWriter.java

@@ -357,8 +357,8 @@ private void writeImpl(Iterator<Product2<K, V>> records) {
     long checkStartTs = System.currentTimeMillis();
     checkAllBufferSpilled();
     checkSentRecordCount(recordCount);
+    checkBlockSendResult(new HashSet<>(blockIds));
     checkSentBlockCount();
-    checkBlockSendResult(blockIds);
     long commitStartTs = System.currentTimeMillis();
     long checkDuration = commitStartTs - checkStartTs;
     if (!isMemoryShuffleEnabled) {
@@ -408,17 +408,26 @@ private void checkSentRecordCount(long recordCount) {
   }
 
   private void checkSentBlockCount() {
-    long tracked = 0;
-    if (serverToPartitionToBlockIds != null) {
-      Set<Long> blockIds = new HashSet<>();
-      for (Map<Integer, Set<Long>> partitionBlockIds : serverToPartitionToBlockIds.values()) {
-        partitionBlockIds.values().forEach(x -> blockIds.addAll(x));
-      }
-      tracked = blockIds.size();
+    long expected = blockIds.size();
+    long bufferManagerTracked = bufferManager.getBlockCount();
+
+    assert serverToPartitionToBlockIds != null;
+    // to filter the multiple replica's duplicate blockIds
+    Set<Long> blockIds = new HashSet<>();
+    for (Map<Integer, Set<Long>> partitionBlockIds : serverToPartitionToBlockIds.values()) {
+      partitionBlockIds.values().forEach(x -> blockIds.addAll(x));
     }
-    if (tracked != bufferManager.getBlockCount()) {
+    long serverTracked = blockIds.size();
+    if (expected != serverTracked || expected != bufferManagerTracked) {
       throw new RssSendFailedException(
-          "Potential block loss may occur when preparing to send blocks for task[" + taskId + "]");
+          "Potential block loss may occur for task["
+              + taskId
+              + "]. BlockId number expected: "
+              + expected
+              + ", serverTracked: "
+              + serverTracked
+              + ", bufferManagerTracked: "
+              + bufferManagerTracked);
     }
   }