(incubator-uniffle) branch master updated: [MINOR] improvement(server): Perf for single Local storage (#2176)

[zuston]

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 b07cc87d7 [MINOR] improvement(server): Perf for single Local storage 
(#2176)
b07cc87d7 is described below

commit b07cc87d7eb9e55ecb30fe729ae9d4603cbc5893
Author: maobaolong <baoloong...@tencent.com>
AuthorDate: Tue Oct 15 11:59:38 2024 +0800

    [MINOR] improvement(server): Perf for single Local storage (#2176)
    
    ### What changes were proposed in this pull request?
    
    Do not select storage when there is only one storage, this can reduce the 
memory cost of `LocalStorageManager#sortedPartitionsOfStorageMap`.
    
    ### Why are the changes needed?
    
    Perf for single Local storage
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    No need.
---
 .../org/apache/uniffle/server/storage/LocalStorageManager.java   | 9 +++++++++
 .../java/org/apache/uniffle/server/ShuffleFlushManagerTest.java  | 3 ++-
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git 
a/server/src/main/java/org/apache/uniffle/server/storage/LocalStorageManager.java
 
b/server/src/main/java/org/apache/uniffle/server/storage/LocalStorageManager.java
index ab5824896..67ace22d8 100644
--- 
a/server/src/main/java/org/apache/uniffle/server/storage/LocalStorageManager.java
+++ 
b/server/src/main/java/org/apache/uniffle/server/storage/LocalStorageManager.java
@@ -204,6 +204,12 @@ public class LocalStorageManager extends 
SingleStorageManager {
 
   @Override
   public Storage selectStorage(ShuffleDataFlushEvent event) {
+    if (localStorages.size() == 1) {
+      if (event.getUnderStorage() == null) {
+        event.setUnderStorage(localStorages.get(0));
+      }
+      return localStorages.get(0);
+    }
     String appId = event.getAppId();
     int shuffleId = event.getShuffleId();
     int partitionId = event.getStartPartition();
@@ -254,6 +260,9 @@ public class LocalStorageManager extends 
SingleStorageManager {
 
   @Override
   public Storage selectStorage(ShuffleDataReadEvent event) {
+    if (localStorages.size() == 1) {
+      return localStorages.get(0);
+    }
     String appId = event.getAppId();
     int shuffleId = event.getShuffleId();
     int partitionId = event.getStartPartition();
diff --git 
a/server/src/test/java/org/apache/uniffle/server/ShuffleFlushManagerTest.java 
b/server/src/test/java/org/apache/uniffle/server/ShuffleFlushManagerTest.java
index 59e92e320..1c47c853f 100644
--- 
a/server/src/test/java/org/apache/uniffle/server/ShuffleFlushManagerTest.java
+++ 
b/server/src/test/java/org/apache/uniffle/server/ShuffleFlushManagerTest.java
@@ -607,7 +607,8 @@ public class ShuffleFlushManagerTest extends HadoopTestBase 
{
 
     ShuffleDataReadEvent shuffle1ReadEvent = new ShuffleDataReadEvent(appId2, 
1, 0, 0);
     ShuffleDataReadEvent shuffle11ReadEvent = new ShuffleDataReadEvent(appId2, 
11, 0, 0);
-    assertNull(storageManager.selectStorage(shuffle1ReadEvent));
+    // As there are only one storage, it should select the only storage anyway
+    assertNotNull(storageManager.selectStorage(shuffle1ReadEvent));
     assertNotNull(storageManager.selectStorage(shuffle11ReadEvent));
 
     storageManager.removeResources(