[#1472][FOLLOWUP] improvement(server): Release memory more accurately when failing to cache shuffle data (#1597)
### What changes were proposed in this pull request?
Release memory more accurately when failing to cache shuffle data.
### Why are the changes needed?
A follow-up PR for: https://github.com/apache/incubator-uniffle/pull/1534.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing UTs.
diff --git a/server/src/main/java/org/apache/uniffle/server/ShuffleServerGrpcService.java b/server/src/main/java/org/apache/uniffle/server/ShuffleServerGrpcService.java
index 695c2af..dcc2717 100644
--- a/server/src/main/java/org/apache/uniffle/server/ShuffleServerGrpcService.java
+++ b/server/src/main/java/org/apache/uniffle/server/ShuffleServerGrpcService.java
@@ -298,11 +298,14 @@
LOG.error(errorMsg);
hasFailureOccurred = true;
break;
+ } finally {
+ if (hasFailureOccurred) {
+ shuffleServer
+ .getShuffleBufferManager()
+ .releaseMemory(spd.getTotalBlockSize(), false, false);
+ }
}
}
- if (hasFailureOccurred) {
- shuffleServer.getShuffleBufferManager().releaseMemory(info.getRequireSize(), false, false);
- }
// since the required buffer id is only used once, the shuffle client would try to require
// another buffer whether
// current connection succeeded or not. Therefore, the preAllocatedBuffer is first get and
