[apache#1628] Avoid exception caused by calling release multiple times

wForget · Jun 24, 2024 · d1be2e4 · d1be2e4
1 parent 1482804
commit d1be2e4
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 2 deletions.
diff --git a/common/src/main/java/org/apache/uniffle/common/ShuffleDataResult.java b/common/src/main/java/org/apache/uniffle/common/ShuffleDataResult.java
@@ -30,7 +30,7 @@
 
 public class ShuffleDataResult {
 
-  private final ManagedBuffer buffer;
+  private volatile ManagedBuffer buffer;
   private final List<BufferSegment> bufferSegments;
 
   public ShuffleDataResult() {
@@ -109,5 +109,6 @@ public void release() {
     if (this.buffer != null) {
       this.buffer.release();
     }
+    this.buffer = null;
   }
 }
diff --git a/common/src/main/java/org/apache/uniffle/common/ShuffleIndexResult.java b/common/src/main/java/org/apache/uniffle/common/ShuffleIndexResult.java
@@ -26,7 +26,7 @@
 import org.apache.uniffle.common.util.ByteBufUtils;
 
 public class ShuffleIndexResult {
-  private final ManagedBuffer buffer;
+  private volatile ManagedBuffer buffer;
   private long dataFileLen;
 
   public ShuffleIndexResult() {
@@ -74,6 +74,7 @@ public void release() {
     if (this.buffer != null) {
       this.buffer.release();
     }
+    this.buffer = null;
   }
 
   public ManagedBuffer getManagedBuffer() {

diff --git a/common/src/test/java/org/apache/uniffle/common/ShuffleDataResultTest.java b/common/src/test/java/org/apache/uniffle/common/ShuffleDataResultTest.java
@@ -22,6 +22,7 @@
 
 import org.junit.jupiter.api.Test;
 
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -39,4 +40,13 @@ public void testEmpty() {
     assertTrue(new ShuffleDataResult(new byte[1], Collections.emptyList()).isEmpty());
     assertFalse(new ShuffleDataResult(new byte[1], segments).isEmpty());
   }
+
+  @Test
+  public void testRelease() {
+    ShuffleDataResult shuffleDataResult = new ShuffleDataResult("test".getBytes(), null);
+    shuffleDataResult.release();
+    // Expect no exception when executing release again
+    assertDoesNotThrow(shuffleDataResult::release);
+    assertTrue(shuffleDataResult.isEmpty());
+  }
 }
diff --git a/common/src/test/java/org/apache/uniffle/common/ShuffleIndexResultTest.java b/common/src/test/java/org/apache/uniffle/common/ShuffleIndexResultTest.java
@@ -19,6 +19,7 @@
 
 import org.junit.jupiter.api.Test;
 
+import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class ShuffleIndexResultTest {
@@ -28,4 +29,13 @@ public void testEmpty() {
     assertTrue(new ShuffleIndexResult().isEmpty());
     assertTrue(new ShuffleIndexResult((byte[]) null, -1).isEmpty());
   }
+
+  @Test
+  public void testRelease() {
+    ShuffleIndexResult shuffleIndexResult = new ShuffleIndexResult("test".getBytes(), -1);
+    shuffleIndexResult.release();
+    // Expect no exception when executing release again
+    assertDoesNotThrow(shuffleIndexResult::release);
+    assertTrue(shuffleIndexResult.isEmpty());
+  }
 }