fix(server): Fix netty memory leak when removeBuffer and cacheShuffle…

…Data happen concurrent (merge request !278)
apache · Oct 10, 2024 · 30229b5 · 30229b5
1 parent ba2302c
commit 30229b5
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 22 deletions.
diff --git a/server/src/main/java/org/apache/uniffle/server/buffer/AbstractShuffleBuffer.java b/server/src/main/java/org/apache/uniffle/server/buffer/AbstractShuffleBuffer.java
@@ -41,8 +41,12 @@ public abstract class AbstractShuffleBuffer implements ShuffleBuffer {
 
   protected long size;
 
+  protected volatile boolean evicted;
+  public static final long BUFFER_EVICTED = -1L;
+
   public AbstractShuffleBuffer() {
     this.size = 0;
+    this.evicted = false;
   }
 
   /** Only for test */

diff --git a/server/src/main/java/org/apache/uniffle/server/buffer/ShuffleBufferManager.java b/server/src/main/java/org/apache/uniffle/server/buffer/ShuffleBufferManager.java
@@ -194,6 +194,9 @@ public StatusCode cacheShuffleData(
 
     ShuffleBuffer buffer = entry.getValue();
     long size = buffer.append(spd);
+    if (size == AbstractShuffleBuffer.BUFFER_EVICTED) {
+      return StatusCode.NO_REGISTER;
+    }
     if (!isPreAllocated) {
       updateUsedMemory(size);
     }

diff --git a/server/src/main/java/org/apache/uniffle/server/buffer/ShuffleBufferWithLinkedList.java b/server/src/main/java/org/apache/uniffle/server/buffer/ShuffleBufferWithLinkedList.java
@@ -48,18 +48,19 @@ public ShuffleBufferWithLinkedList() {
   }
 
   @Override
-  public long append(ShufflePartitionedData data) {
-    long size = 0;
+  public synchronized long append(ShufflePartitionedData data) {
+    if (evicted) {
+      return BUFFER_EVICTED;
+    }
+    long currentSize = 0;
 
-    synchronized (this) {
-      for (ShufflePartitionedBlock block : data.getBlockList()) {
-        blocks.add(block);
-        size += block.getSize();
-      }
-      this.size += size;
+    for (ShufflePartitionedBlock block : data.getBlockList()) {
+      blocks.add(block);
+      currentSize += block.getSize();
     }
+    this.size += currentSize;
 
-    return size;
+    return currentSize;
   }
 
   @Override
@@ -110,7 +111,8 @@ public int getBlockCount() {
   }
 
   @Override
-  public void release() {
+  public synchronized void release() {
+    evicted = true;
     blocks.forEach(spb -> spb.getData().release());
   }
 

diff --git a/server/src/main/java/org/apache/uniffle/server/buffer/ShuffleBufferWithSkipList.java b/server/src/main/java/org/apache/uniffle/server/buffer/ShuffleBufferWithSkipList.java
@@ -57,19 +57,20 @@ private ConcurrentSkipListMap<Long, ShufflePartitionedBlock> newConcurrentSkipLi
   }
 
   @Override
-  public long append(ShufflePartitionedData data) {
-    long size = 0;
-
-    synchronized (this) {
-      for (ShufflePartitionedBlock block : data.getBlockList()) {
-        blocksMap.put(block.getBlockId(), block);
-        blockCount++;
-        size += block.getSize();
-      }
-      this.size += size;
+  public synchronized long append(ShufflePartitionedData data) {
+    if (evicted) {
+      return BUFFER_EVICTED;
+    }
+    long currentSize = 0;
+
+    for (ShufflePartitionedBlock block : data.getBlockList()) {
+      blocksMap.put(block.getBlockId(), block);
+      blockCount++;
+      currentSize += block.getSize();
     }
+    this.size += currentSize;
 
-    return size;
+    return currentSize;
   }
 
   @Override
@@ -111,7 +112,8 @@ public int getBlockCount() {
   }
 
   @Override
-  public void release() {
+  public synchronized void release() {
+    evicted = true;
     blocksMap.values().forEach(spb -> spb.getData().release());
   }