[incubator-uniffle] 03/04: Fix potenial race condition when registering remote storage info (#481)

[ck...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

ckj pushed a commit to branch branch-0.6
in repository https://gitbox.apache.org/repos/asf/incubator-uniffle.git

commit 9f3d0743363cd2d2faf414ac3a389328c2a27168
Author: Junfan Zhang <zu...@apache.org>
AuthorDate: Mon Jan 16 13:56:37 2023 +0800

    Fix potenial race condition when registering remote storage info (#481)
    
    
    ### What changes were proposed in this pull request?
    
    1. Use the concurrentHashMap's computeIfAbsent to keep thread safe
    
    ### Why are the changes needed?
    1. To fix potential race condition when registering remote storage info
    
    ### Does this PR introduce _any_ user-facing change?
    No
    
    ### How was this patch tested?
    Existing UTs
---
 .../uniffle/server/storage/HdfsStorageManager.java      | 17 +++++++----------
 1 file changed, 7 insertions(+), 10 deletions(-)

diff --git a/server/src/main/java/org/apache/uniffle/server/storage/HdfsStorageManager.java b/server/src/main/java/org/apache/uniffle/server/storage/HdfsStorageManager.java
index 7c0c0dd3..72e622ce 100644
--- a/server/src/main/java/org/apache/uniffle/server/storage/HdfsStorageManager.java
+++ b/server/src/main/java/org/apache/uniffle/server/storage/HdfsStorageManager.java
@@ -95,23 +95,20 @@ public class HdfsStorageManager extends SingleStorageManager {
   @Override
   public void registerRemoteStorage(String appId, RemoteStorageInfo remoteStorageInfo) {
     String remoteStorage = remoteStorageInfo.getPath();
-    Map<String, String> remoteStorageConf = remoteStorageInfo.getConfItems();
-    if (!pathToStorages.containsKey(remoteStorage)) {
+    pathToStorages.computeIfAbsent(remoteStorage, key -> {
+      Map<String, String> remoteStorageConf = remoteStorageInfo.getConfItems();
       Configuration remoteStorageHadoopConf = new Configuration(hadoopConf);
       if (remoteStorageConf != null && remoteStorageConf.size() > 0) {
         for (Map.Entry<String, String> entry : remoteStorageConf.entrySet()) {
           remoteStorageHadoopConf.setStrings(entry.getKey(), entry.getValue());
         }
       }
-      pathToStorages.putIfAbsent(remoteStorage, new HdfsStorage(remoteStorage, remoteStorageHadoopConf));
-      // registerRemoteStorage may be called in different threads,
-      // make sure metrics won't be created duplicated
-      // there shouldn't have performance issue because
-      // it will be called only few times according to the number of remote storage
-      String storageHost = pathToStorages.get(remoteStorage).getStorageHost();
+      HdfsStorage hdfsStorage = new HdfsStorage(remoteStorage, remoteStorageHadoopConf);
+      String storageHost = hdfsStorage.getStorageHost();
       ShuffleServerMetrics.addDynamicCounterForRemoteStorage(storageHost);
-    }
-    appIdToStorages.putIfAbsent(appId, pathToStorages.get(remoteStorage));
+      return hdfsStorage;
+    });
+    appIdToStorages.computeIfAbsent(appId, key -> pathToStorages.get(remoteStorage));
   }
 
   public HdfsStorage getStorageByAppId(String appId) {