You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@uniffle.apache.org by zu...@apache.org on 2026/10/12 07:58:39 UTC
(uniffle) branch master updated: [#2812] fix(spark): Double memory release after shuffle write failure (#2813)
This is an automated email from the ASF dual-hosted git repository.
zuston pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/uniffle.git
The following commit(s) were added to refs/heads/master by this push:
new e78d6466c [#2812] fix(spark): Double memory release after shuffle write failure (#2813)
e78d6466c is described below
commit e78d6466c7b7ff70bf60f6c0133aad7150e598c6
Author: Junfan Zhang <zu...@apache.org>
AuthorDate: Mon Oct 12 15:58:33 2026 +0800
[#2812] fix(spark): Double memory release after shuffle write failure (#2813)
### What changes were proposed in this pull request?
To fix the double memory release after shuffle write failure.
### Why are the changes needed?
Fix: #2812
After a Spark shuffle write fails, `RssShuffleWriter.stop()` calls `WriteBufferManager.freeAllMemory()`. Pending asynchronous block completion callbacks can then release the same memory again, producing repeated warnings:
```text
Internal error: release called on 3072 bytes but task only has 0 bytes of memory from the on-heap execution pool
```
Memory cleanup should be idempotent and coordinated with asynchronous callbacks, while still releasing the blocks' data buffers.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
unit tests
---
.../spark/shuffle/writer/WriteBufferManager.java | 18 ++-
.../shuffle/writer/WriteBufferManagerTest.java | 132 ++++++++++++++++++---
2 files changed, 133 insertions(+), 17 deletions(-)
diff --git a/client-spark/common/src/main/java/org/apache/spark/shuffle/writer/WriteBufferManager.java b/client-spark/common/src/main/java/org/apache/spark/shuffle/writer/WriteBufferManager.java
index 651b0b47e..ed3b342f9 100644
--- a/client-spark/common/src/main/java/org/apache/spark/shuffle/writer/WriteBufferManager.java
+++ b/client-spark/common/src/main/java/org/apache/spark/shuffle/writer/WriteBufferManager.java
@@ -659,6 +659,12 @@ public class WriteBufferManager extends MemoryConsumer {
}
}
+ // Spark 2 declares MemoryConsumer.getUsed() as protected.
+ @Override
+ public long getUsed() {
+ return super.getUsed();
+ }
+
@VisibleForTesting
protected long getAllocatedBytes() {
return allocatedBytes.get();
@@ -687,15 +693,21 @@ public class WriteBufferManager extends MemoryConsumer {
return blockCounter.get();
}
- public void freeAllocatedMemory(long freeMemory) {
+ public synchronized void freeAllocatedMemory(long freeMemory) {
+ // Completion callbacks may arrive after freeAllMemory has released the task's memory.
+ if (allocatedBytes.get() == 0) {
+ return;
+ }
freeMemory(freeMemory);
allocatedBytes.addAndGet(-freeMemory);
usedBytes.addAndGet(-freeMemory);
inSendListBytes.addAndGet(-freeMemory);
}
- public void freeAllMemory() {
- long memory = allocatedBytes.get();
+ public synchronized void freeAllMemory() {
+ long memory = allocatedBytes.getAndSet(0);
+ usedBytes.set(0);
+ inSendListBytes.set(0);
if (memory > 0) {
freeMemory(memory);
}
diff --git a/client-spark/common/src/test/java/org/apache/spark/shuffle/writer/WriteBufferManagerTest.java b/client-spark/common/src/test/java/org/apache/spark/shuffle/writer/WriteBufferManagerTest.java
index 4601b3779..8603e2080 100644
--- a/client-spark/common/src/test/java/org/apache/spark/shuffle/writer/WriteBufferManagerTest.java
+++ b/client-spark/common/src/test/java/org/apache/spark/shuffle/writer/WriteBufferManagerTest.java
@@ -23,6 +23,7 @@ import java.util.Arrays;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import java.util.stream.Stream;
@@ -43,6 +44,8 @@ import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.ValueSource;
+import org.mockito.ArgumentCaptor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -56,33 +59,37 @@ import static org.apache.spark.shuffle.RssSparkConfig.toSparkConfKey;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.same;
+import static org.mockito.Mockito.atLeastOnce;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
public class WriteBufferManagerTest {
private WriteBufferManager createManager(SparkConf conf) {
- Serializer kryoSerializer = new KryoSerializer(conf);
TaskMemoryManager mockTaskMemoryManager = mock(TaskMemoryManager.class);
-
- BufferManagerOptions bufferOptions = new BufferManagerOptions(conf);
- WriteBufferManager wbm =
- new WriteBufferManager(
- 0,
- 0,
- bufferOptions,
- kryoSerializer,
- Maps.newHashMap(),
- mockTaskMemoryManager,
- new ShuffleWriteMetrics(),
- RssSparkConfig.toRssConf(conf));
- WriteBufferManager spyManager = spy(wbm);
+ WriteBufferManager spyManager = spy(createManager(conf, mockTaskMemoryManager));
doReturn(512L).when(spyManager).acquireMemory(anyLong());
return spyManager;
}
+ private WriteBufferManager createManager(SparkConf conf, TaskMemoryManager taskMemoryManager) {
+ Serializer kryoSerializer = new KryoSerializer(conf);
+ BufferManagerOptions bufferOptions = new BufferManagerOptions(conf);
+ return new WriteBufferManager(
+ 0,
+ 0,
+ bufferOptions,
+ kryoSerializer,
+ Maps.newHashMap(),
+ taskMemoryManager,
+ new ShuffleWriteMetrics(),
+ RssSparkConfig.toRssConf(conf));
+ }
+
private SparkConf getConf() {
SparkConf conf = new SparkConf(false);
conf.set(toSparkConfKey(RssClientConf.RSS_WRITER_BUFFER_SIZE), "64")
@@ -199,6 +206,103 @@ public class WriteBufferManagerTest {
assertEquals(96, wbm.getInSendListBytes());
}
+ @ParameterizedTest
+ @ValueSource(booleans = {true, false})
+ public void freeAllMemoryWithPendingCallbacksTest(boolean cleanupFirst) {
+ SparkConf conf = getConf();
+ TaskMemoryManager taskMemoryManager = mock(TaskMemoryManager.class);
+ WriteBufferManager wbm = createManager(conf, taskMemoryManager);
+ when(taskMemoryManager.acquireExecutionMemory(anyLong(), same(wbm))).thenReturn(512L);
+
+ wbm.addRecord(0, "Key", "Value");
+ wbm.addRecord(1, "Key", "Value");
+ List<ShuffleBlockInfo> blocks = wbm.clear();
+ wbm.buildBlockEvents(blocks);
+ wbm.addRecord(2, "Key", "Value");
+ assertEquals(512, wbm.getAllocatedBytes());
+ assertEquals(512, wbm.getUsed());
+ assertEquals(96, wbm.getUsedBytes());
+ assertEquals(64, wbm.getInSendListBytes());
+ blocks.forEach(block -> assertEquals(1, block.getData().refCnt()));
+
+ if (!cleanupFirst) {
+ blocks.get(0).executeCompletionCallback(true);
+ assertEquals(480, wbm.getAllocatedBytes());
+ assertEquals(64, wbm.getUsedBytes());
+ assertEquals(32, wbm.getInSendListBytes());
+ }
+ wbm.freeAllMemory();
+ wbm.freeAllMemory();
+ if (cleanupFirst) {
+ blocks.get(0).executeCompletionCallback(true);
+ }
+ blocks.get(1).executeCompletionCallback(false);
+ wbm.freeAllMemory();
+
+ assertEquals(0, wbm.getAllocatedBytes());
+ assertEquals(0, wbm.getUsed());
+ assertEquals(0, wbm.getUsedBytes());
+ assertEquals(0, wbm.getInSendListBytes());
+ blocks.forEach(block -> assertEquals(0, block.getData().refCnt()));
+ ArgumentCaptor<Long> releasedBytes = ArgumentCaptor.forClass(Long.class);
+ verify(taskMemoryManager, atLeastOnce())
+ .releaseExecutionMemory(releasedBytes.capture(), same(wbm));
+ assertEquals(512, releasedBytes.getAllValues().stream().mapToLong(Long::longValue).sum());
+ }
+
+ @Test
+ public void freeAllMemoryWithConcurrentCallbacksTest() throws Exception {
+ SparkConf conf = getConf();
+ conf.set(toSparkConfKey(RssSparkConfig.RSS_WRITER_BUFFER_SPILL_SIZE), "1024");
+ TaskMemoryManager taskMemoryManager = mock(TaskMemoryManager.class);
+ WriteBufferManager wbm = createManager(conf, taskMemoryManager);
+ when(taskMemoryManager.acquireExecutionMemory(anyLong(), same(wbm))).thenReturn(512L);
+
+ for (int partition = 0; partition < 16; partition++) {
+ wbm.addRecord(partition, "Key", "Value");
+ }
+ List<ShuffleBlockInfo> blocks = wbm.clear();
+ wbm.buildBlockEvents(blocks);
+ CountDownLatch start = new CountDownLatch(1);
+ List<CompletableFuture<Void>> futures = new ArrayList<>();
+ for (ShuffleBlockInfo block : blocks) {
+ futures.add(
+ CompletableFuture.runAsync(
+ () -> {
+ awaitLatch(start);
+ block.executeCompletionCallback(false);
+ }));
+ }
+ futures.add(
+ CompletableFuture.runAsync(
+ () -> {
+ awaitLatch(start);
+ wbm.freeAllMemory();
+ wbm.freeAllMemory();
+ }));
+ start.countDown();
+ CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).get(10, TimeUnit.SECONDS);
+
+ assertEquals(0, wbm.getAllocatedBytes());
+ assertEquals(0, wbm.getUsed());
+ assertEquals(0, wbm.getUsedBytes());
+ assertEquals(0, wbm.getInSendListBytes());
+ blocks.forEach(block -> assertEquals(0, block.getData().refCnt()));
+ ArgumentCaptor<Long> releasedBytes = ArgumentCaptor.forClass(Long.class);
+ verify(taskMemoryManager, atLeastOnce())
+ .releaseExecutionMemory(releasedBytes.capture(), same(wbm));
+ assertEquals(512, releasedBytes.getAllValues().stream().mapToLong(Long::longValue).sum());
+ }
+
+ private static void awaitLatch(CountDownLatch latch) {
+ try {
+ assertTrue(latch.await(10, TimeUnit.SECONDS));
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new AssertionError(e);
+ }
+ }
+
@Test
public void addHugeRecordTest() {
SparkConf conf = getConf();