fix getLocalShuffleIndex

zhengchenyu · Aug 17, 2023 · f12143c · f12143c
1 parent fe28c86
commit f12143c
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/common/src/main/java/org/apache/uniffle/common/ShuffleIndexResult.java b/common/src/main/java/org/apache/uniffle/common/ShuffleIndexResult.java
@@ -69,4 +69,10 @@ public long getDataFileLen() {
  public boolean isEmpty() {
  return buffer == null || buffer.size() == 0;
  }
+
+ public void release() {
+ if (this.buffer != null) {
+ this.buffer.release();
+ }
+ }
 }
diff --git a/server/src/main/java/org/apache/uniffle/server/ShuffleServerGrpcService.java b/server/src/main/java/org/apache/uniffle/server/ShuffleServerGrpcService.java
@@ -709,9 +709,10 @@ public void getLocalShuffleIndex(
  .getShuffleServerConf()
  .getLong(ShuffleServerConf.SERVER_SHUFFLE_INDEX_SIZE_HINT);
  if (shuffleServer.getShuffleBufferManager().requireReadMemoryWithRetry(assumedFileSize)) {
+ ShuffleIndexResult shuffleIndexResult = null;
  try {
  long start = System.currentTimeMillis();
- ShuffleIndexResult shuffleIndexResult =
+ shuffleIndexResult =
  shuffleServer
  .getShuffleTaskManager()
  .getShuffleIndex(appId, shuffleId, partitionId, partitionNumPerRange, partitionNum);
@@ -747,6 +748,9 @@ public void getLocalShuffleIndex(
  .setRetMsg(msg)
  .build();
  } finally {
+ if (shuffleIndexResult != null) {
+ shuffleIndexResult.release();
+ }
  shuffleServer.getShuffleBufferManager().releaseReadMemory(assumedFileSize);
  }
  } else {