(incubator-uniffle) branch master updated: [#1472][FOLLOWUP] improvement(server): Release memory more accurately when failing to cache shuffle data (#1597)

[zu...@apache.org]

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/incubator-uniffle.git


The following commit(s) were added to refs/heads/master by this push:
     new 220d9cbe5 [#1472][FOLLOWUP] improvement(server): Release memory more accurately when failing to cache shuffle data (#1597)
220d9cbe5 is described below

commit 220d9cbe5a573f9ca35a5e6eaad29e5648736ce1
Author: RickyMa <ri...@tencent.com>
AuthorDate: Mon Mar 25 09:57:25 2024 +0800

    [#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.
---
 .../java/org/apache/uniffle/server/ShuffleServerGrpcService.java | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

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 695c2afc6..dcc2717c5 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 @@ public class ShuffleServerGrpcService extends ShuffleServerImplBase {
           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