From d1be2e4c6d8b4273a9144d65d6002a32cc0ef6b8 Mon Sep 17 00:00:00 2001
From: wforget <643348094@qq.com>
Date: Mon, 24 Jun 2024 15:35:50 +0800
Subject: [PATCH] [#1628] Avoid exception caused by calling release multiple
 times

---
 .../org/apache/uniffle/common/ShuffleDataResult.java   |  3 ++-
 .../org/apache/uniffle/common/ShuffleIndexResult.java  |  3 ++-
 .../apache/uniffle/common/ShuffleDataResultTest.java   | 10 ++++++++++
 .../apache/uniffle/common/ShuffleIndexResultTest.java  | 10 ++++++++++
 4 files changed, 24 insertions(+), 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
index 00867e7006..d4e7d06807 100644
--- 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
index 71bb3df393..0700b07048 100644
--- 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
index 5c97eb69b1..756ad0cfb5 100644
--- 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
index 3637029f8f..696fdfc482 100644
--- 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());
+  }
 }