[#2411] fix(spark): Spill memory corresponding to successfully sent blocks (#2415)

### What changes were proposed in this pull request?
As title

### Why are the changes needed?
Before this pr, spark client spill more memory than actually did

### Does this PR introduce _any_ user-facing change?
No.

### How was this patch tested?
UT

Author: summaryzb

client-spark/common/src/main/java/org/apache/spark/shuffle/writer/DataPusher.java

@@ -100,10 +100,7 @@ public CompletableFuture<Long> send(AddBlockEvent event) {
                 putFailedBlockSendTracker(
                     taskToFailedBlockSendTracker, taskId, result.getFailedBlockSendTracker());
               } finally {
-                Set<Long> succeedBlockIds =
-                    result.getSuccessBlockIds() == null
-                        ? Collections.emptySet()
-                        : result.getSuccessBlockIds();
+                Set<Long> succeedBlockIds = getSucceedBlockIds(result);
                 for (ShuffleBlockInfo block : shuffleBlockInfoList) {
                   block.executeCompletionCallback(succeedBlockIds.contains(block.getBlockId()));
                 }
@@ -114,7 +111,9 @@ public CompletableFuture<Long> send(AddBlockEvent event) {
                   runnable.run();
                 }
               }
+              Set<Long> succeedBlockIds = getSucceedBlockIds(result);
               return shuffleBlockInfoList.stream()
+                  .filter(x -> succeedBlockIds.contains(x.getBlockId()))
                   .map(x -> x.getFreeMemory())
                   .reduce((a, b) -> a + b)
                   .get();
@@ -127,6 +126,13 @@ public CompletableFuture<Long> send(AddBlockEvent event) {
             });
   }
 
+  private Set<Long> getSucceedBlockIds(SendShuffleDataResult result) {
+    if (result == null || result.getSuccessBlockIds() == null) {
+      return Collections.emptySet();
+    }
+    return result.getSuccessBlockIds();
+  }
+
   private synchronized void putBlockId(
       Map<String, Set<Long>> taskToBlockIds, String taskAttemptId, Set<Long> blockIds) {
     if (blockIds == null || blockIds.isEmpty()) {

client-spark/common/src/test/java/org/apache/spark/shuffle/writer/WriteBufferManagerTest.java

@@ -487,7 +487,12 @@ public void spillPartial() {
           List<AddBlockEvent> events = wbm.buildBlockEvents(blocks);
           for (AddBlockEvent event : events) {
             event.getProcessedCallbackChain().stream().forEach(x -> x.run());
-            sum += event.getShuffleDataInfoList().stream().mapToLong(x -> x.getFreeMemory()).sum();
+            // simulate: the block for partition 2 send failed
+            sum +=
+                event.getShuffleDataInfoList().stream()
+                    .filter(x -> x.getPartitionId() <= 1)
+                    .mapToLong(x -> x.getFreeMemory())
+                    .sum();
           }
           return Arrays.asList(CompletableFuture.completedFuture(sum));
         };
@@ -502,10 +507,11 @@ public void spillPartial() {
     wbm.addRecord(1, testKey, testValue);
     wbm.addRecord(1, testKey, testValue);
     wbm.addRecord(1, testKey, testValue);
+    wbm.addRecord(2, testKey, testValue);
 
     long releasedSize = wbm.spill(1000, wbm);
     assertEquals(64, releasedSize);
-    assertEquals(96, wbm.getUsedBytes());
+    assertEquals(128, wbm.getUsedBytes());
     assertEquals(0, wbm.getBuffers().keySet().toArray()[0]);
   }