[CELEBORN-1544][FOLLOWUP] ShuffleWriter needs to catch exception and …

…call abort to avoid memory leaks

### What changes were proposed in this pull request?
This PR aims to fix a possible memory leak in ShuffleWriter.

Introduce a private abort method, which can be called to release memory when an exception occurs.

### Why are the changes needed?
#2661 Call the close method in the finally block, but the close method has `shuffleClient.mapperEnd`, which is dangerous for incomplete tasks, and the data may be inaccurate.

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

### How was this patch tested?
GA

Closes #2663 from cxzl25/CELEBORN-1544-followup.

Authored-by: sychen <[email protected]>
Signed-off-by: zky.zhoukeyong <[email protected]>
(cherry picked from commit bc3bd46)
Signed-off-by: zky.zhoukeyong <[email protected]>

client-spark/common/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedPusher.java

@@ -501,13 +501,15 @@ public long getPushSortMemoryThreshold() {
     return this.pushSortMemoryThreshold;
   }
 
-  public void close() throws IOException {
+  public void close(boolean throwTaskKilledOnInterruption) throws IOException {
     cleanupResources();
     try {
       dataPusher.waitOnTermination();
       sendBufferPool.returnPushTaskQueue(dataPusher.getIdleQueue());
     } catch (InterruptedException e) {
-      TaskInterruptedHelper.throwTaskKillException();
+      if (throwTaskKilledOnInterruption) {
+        TaskInterruptedHelper.throwTaskKillException();
+      }
     }
   }
 

client-spark/common/src/test/java/org/apache/spark/shuffle/celeborn/SortBasedPusherSuiteJ.java

@@ -127,7 +127,7 @@ public void testMemoryUsage() throws Exception {
         !pusher.insertRecord(
             row5k.getBaseObject(), row5k.getBaseOffset(), row5k.getSizeInBytes(), 0, true));
 
-    pusher.close();
+    pusher.close(true);
 
     assertEquals(taskContext.taskMetrics().memoryBytesSpilled(), 2097152);
   }

client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java

@@ -162,6 +162,7 @@ public HashBasedShuffleWriter(
 
   @Override
   public void write(scala.collection.Iterator<Product2<K, V>> records) throws IOException {
+    boolean needCleanupPusher = true;
     try {
       if (canUseFastWrite()) {
         fastWrite0(records);
@@ -174,13 +175,13 @@ public void write(scala.collection.Iterator<Product2<K, V>> records) throws IOEx
       } else {
         write0(records);
       }
+      close();
+      needCleanupPusher = false;
     } catch (InterruptedException e) {
       TaskInterruptedHelper.throwTaskKillException();
     } finally {
-      try {
-        close();
-      } catch (InterruptedException e) {
-        TaskInterruptedHelper.throwTaskKillException();
+      if (needCleanupPusher) {
+        cleanupPusher();
       }
     }
   }
@@ -319,6 +320,15 @@ private void flushSendBuffer(int partitionId, byte[] buffer, int size)
     writeMetrics.incWriteTime(System.nanoTime() - start);
   }
 
+  private void cleanupPusher() throws IOException {
+    try {
+      dataPusher.waitOnTermination();
+      sendBufferPool.returnPushTaskQueue(dataPusher.getIdleQueue());
+    } catch (InterruptedException e) {
+      TaskInterruptedHelper.throwTaskKillException();
+    }
+  }
+
   private void close() throws IOException, InterruptedException {
     // here we wait for all the in-flight batches to return which sent by dataPusher thread
     dataPusher.waitOnTermination();

client-spark/spark-2/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java

@@ -143,6 +143,7 @@ public SortBasedShuffleWriter(
 
   @Override
   public void write(scala.collection.Iterator<Product2<K, V>> records) throws IOException {
+    boolean needCleanupPusher = true;
     try {
       if (canUseFastWrite()) {
         fastWrite0(records);
@@ -155,8 +156,12 @@ public void write(scala.collection.Iterator<Product2<K, V>> records) throws IOEx
       } else {
         write0(records);
       }
-    } finally {
       close();
+      needCleanupPusher = false;
+    } finally {
+      if (needCleanupPusher) {
+        cleanupPusher();
+      }
     }
   }
 
@@ -291,11 +296,17 @@ private void pushGiantRecord(int partitionId, byte[] buffer, int numBytes) throw
     writeMetrics.incBytesWritten(bytesWritten);
   }
 
+  private void cleanupPusher() throws IOException {
+    if (pusher != null) {
+      pusher.close(false);
+    }
+  }
+
   private void close() throws IOException {
     logger.info("Memory used {}", Utils.bytesToString(pusher.getUsed()));
     long pushStartTime = System.nanoTime();
     pusher.pushData(false);
-    pusher.close();
+    pusher.close(true);
     writeMetrics.incWriteTime(System.nanoTime() - pushStartTime);
 
     shuffleClient.pushMergedData(shuffleId, mapId, taskContext.attemptNumber());

client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/HashBasedShuffleWriter.java

@@ -158,6 +158,7 @@ public HashBasedShuffleWriter(
 
   @Override
   public void write(scala.collection.Iterator<Product2<K, V>> records) throws IOException {
+    boolean needCleanupPusher = true;
     try {
       if (canUseFastWrite()) {
         fastWrite0(records);
@@ -170,13 +171,13 @@ public void write(scala.collection.Iterator<Product2<K, V>> records) throws IOEx
       } else {
         write0(records);
       }
+      close();
+      needCleanupPusher = false;
     } catch (InterruptedException e) {
       TaskInterruptedHelper.throwTaskKillException();
     } finally {
-      try {
-        close();
-      } catch (InterruptedException e) {
-        TaskInterruptedHelper.throwTaskKillException();
+      if (needCleanupPusher) {
+        cleanupPusher();
       }
     }
   }
@@ -353,6 +354,15 @@ protected void mergeData(int partitionId, byte[] buffer, int offset, int length)
     writeMetrics.incBytesWritten(bytesWritten);
   }
 
+  private void cleanupPusher() throws IOException {
+    try {
+      dataPusher.waitOnTermination();
+      sendBufferPool.returnPushTaskQueue(dataPusher.getIdleQueue());
+    } catch (InterruptedException e) {
+      TaskInterruptedHelper.throwTaskKillException();
+    }
+  }
+
   private void close() throws IOException, InterruptedException {
     // here we wait for all the in-flight batches to return which sent by dataPusher thread
     long pushMergedDataTime = System.nanoTime();

client-spark/spark-3/src/main/java/org/apache/spark/shuffle/celeborn/SortBasedShuffleWriter.java

@@ -226,10 +226,15 @@ void doWrite(scala.collection.Iterator<Product2<K, V>> records) throws IOExcepti
 
   @Override
   public void write(scala.collection.Iterator<Product2<K, V>> records) throws IOException {
+    boolean needCleanupPusher = true;
     try {
       doWrite(records);
-    } finally {
       close();
+      needCleanupPusher = false;
+    } finally {
+      if (needCleanupPusher) {
+        cleanupPusher();
+      }
     }
   }
 
@@ -354,11 +359,17 @@ private void pushGiantRecord(int partitionId, byte[] buffer, int numBytes) throw
     writeMetrics.incBytesWritten(bytesWritten);
   }
 
+  private void cleanupPusher() throws IOException {
+    if (pusher != null) {
+      pusher.close(false);
+    }
+  }
+
   private void close() throws IOException {
     logger.info("Memory used {}", Utils.bytesToString(pusher.getUsed()));
     long pushStartTime = System.nanoTime();
     pusher.pushData(false);
-    pusher.close();
+    pusher.close(true);
 
     shuffleClient.pushMergedData(shuffleId, mapId, taskContext.attemptNumber());
     writeMetrics.incWriteTime(System.nanoTime() - pushStartTime);