commit 36759f382e9c184c9ec727b2d04a04c90ca7e16f
author Junfan Zhang <zuston@apache.org> Tue Nov 29 14:04:26 2022 +0800
committer Kaijie Chen <ckj@apache.org> Wed Nov 30 17:50:47 2022 +0800
tree 39ed0a0dd668bcc60789c19e6955eee8b41a62fe
parent 753f7b3044f35042f5a0718127f9148db80accbb

[BUG] Potenial memory leak when encountering disk unhealthy (#370)

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

Fix potential memory leak when encountering disk unhealthy

### Why are the changes needed?

When encountering disk unhealthy and exceed the timeout of pendingDataShuffleFlushEvent, it will release memory. But in current codebase, it wont release the data reference and cause the memory leak.

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

No

### How was this patch tested?
No need.

server/src/main/java/org/apache/uniffle/server/ShuffleFlushManager.java

diff --git a/server/src/main/java/org/apache/uniffle/server/ShuffleFlushManager.java b/server/src/main/java/org/apache/uniffle/server/ShuffleFlushManager.java
index daae92d..96f8929 100644
--- a/server/src/main/java/org/apache/uniffle/server/ShuffleFlushManager.java
+++ b/server/src/main/java/org/apache/uniffle/server/ShuffleFlushManager.java
@@ -208,12 +208,8 @@
       // just log the error, don't throw the exception and stop the flush thread
       LOG.error("Exception happened when process flush shuffle data for " + event, e);
     } finally {
-      ShuffleBuffer shuffleBuffer = event.getShuffleBuffer();
-      if (shuffleBuffer != null) {
-        shuffleBuffer.clearInFlushBuffer(event.getEventId());
-      }
+      cleanupFlushEventData(event);
       if (shuffleServer != null) {
-        shuffleServer.getShuffleBufferManager().releaseMemory(event.getSize(), true, false);
         long duration = System.currentTimeMillis() - start;
         if (writeSuccess) {
           LOG.debug("Flush to file success in " + duration + " ms and release " + event.getSize() + " bytes");
@@ -310,14 +306,22 @@
     addPendingEventsInternal(event);
   }
 
-  private void dropPendingEvent(PendingShuffleFlushEvent event) {
-    ShuffleServerMetrics.counterTotalDroppedEventNum.inc();
+  private void cleanupFlushEventData(ShuffleDataFlushEvent event) {
+    ShuffleBuffer shuffleBuffer = event.getShuffleBuffer();
+    if (shuffleBuffer != null) {
+      shuffleBuffer.clearInFlushBuffer(event.getEventId());
+    }
     if (shuffleServer != null) {
       shuffleServer.getShuffleBufferManager().releaseMemory(
-          event.getEvent().getSize(), true, false);
+          event.getSize(), true, false);
     }
   }
 
+  private void dropPendingEvent(PendingShuffleFlushEvent event) {
+    ShuffleServerMetrics.counterTotalDroppedEventNum.inc();
+    cleanupFlushEventData(event.getEvent());
+  }
+
   @VisibleForTesting
   void addPendingEvents(ShuffleDataFlushEvent event) {
     addPendingEventsInternal(new PendingShuffleFlushEvent(event));