Fix.

apache · Jul 23, 2023 · c58a624 · c58a624
1 parent 50f7132
commit c58a624
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 2 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
@@ -386,9 +386,17 @@ public List<AddBlockEvent> buildBlockEvents(List<ShuffleBlockInfo> shuffleBlockI
               + " bytes");
       // Use final temporary variables for closures
       final long _memoryUsed = memoryUsed;
+      final List<ShuffleBlockInfo> finalShuffleBlockInfosPerEvent = shuffleBlockInfoList;
       events.add(
           new AddBlockEvent(
-              taskId, shuffleBlockInfosPerEvent, () -> freeAllocatedMemory(_memoryUsed)));
+              taskId,
+              shuffleBlockInfosPerEvent,
+              () -> {
+                freeAllocatedMemory(_memoryUsed);
+                for (ShuffleBlockInfo shuffleBlockInfo : finalShuffleBlockInfosPerEvent) {
+                  shuffleBlockInfo.getData().release();
+                }
+              }));
     }
     return events;
   }

diff --git a/common/src/main/java/org/apache/uniffle/common/netty/protocol/Encoders.java b/common/src/main/java/org/apache/uniffle/common/netty/protocol/Encoders.java
@@ -44,7 +44,6 @@ public static void encodeShuffleBlockInfo(ShuffleBlockInfo shuffleBlockInfo, Byt
     byteBuf.writeLong(shuffleBlockInfo.getTaskAttemptId());
     // todo: avoid copy
     shuffleBlockInfo.copyDataTo(byteBuf);
-    shuffleBlockInfo.getData().release();
     List<ShuffleServerInfo> shuffleServerInfoList = shuffleBlockInfo.getShuffleServerInfos();
     byteBuf.writeInt(shuffleServerInfoList.size());
     for (ShuffleServerInfo shuffleServerInfo : shuffleServerInfoList) {

diff --git a/...t/spark3/src/test/java/org/apache/uniffle/test/ContinuousSelectPartitionStrategyTest.java b/...t/spark3/src/test/java/org/apache/uniffle/test/ContinuousSelectPartitionStrategyTest.java
@@ -95,6 +95,7 @@ private static void createShuffleServers() throws Exception {
       serverConf.setString("rss.server.hadoop.dfs.replication", "2");
       serverConf.setLong("rss.server.disk.capacity", 10L * 1024L * 1024L * 1024L);
       serverConf.setBoolean("rss.server.health.check.enable", false);
+      serverConf.setString("rss.server.tags", "GRPC,GRPC_NETTY");
       createMockedShuffleServer(serverConf);
     }
     enableRecordGetShuffleResult();

diff --git a/...n-test/spark3/src/test/java/org/apache/uniffle/test/GetShuffleReportForMultiPartTest.java b/...n-test/spark3/src/test/java/org/apache/uniffle/test/GetShuffleReportForMultiPartTest.java
@@ -99,6 +99,7 @@ private static void createShuffleServers() throws Exception {
       serverConf.setString("rss.server.hadoop.dfs.replication", "2");
       serverConf.setLong("rss.server.disk.capacity", 10L * 1024L * 1024L * 1024L);
       serverConf.setBoolean("rss.server.health.check.enable", false);
+      serverConf.setString("rss.server.tags", "GRPC,GRPC_NETTY");
       createMockedShuffleServer(serverConf);
     }
     enableRecordGetShuffleResult();