[#493][FOLLOWUP] improvement: replace putIfAbsent to avoid performance loss (#2444)

### What changes were proposed in this pull request?

Replace putIfAbsent to avoid performance loss

### Why are the changes needed?

Fix: #493

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

CI

Author: xianjingfeng

client-spark/spark2/src/main/java/org/apache/spark/shuffle/RssShuffleManager.java

@@ -115,7 +115,8 @@ public <K, V, C> ShuffleHandle registerShuffle(
 
     if (dependency.partitioner().numPartitions() == 0) {
       shuffleIdToPartitionNum.putIfAbsent(shuffleId, 0);
-      shuffleIdToNumMapTasks.putIfAbsent(shuffleId, dependency.rdd().partitions().length);
+      shuffleIdToNumMapTasks.computeIfAbsent(
+          shuffleId, key -> dependency.rdd().partitions().length);
       LOG.info(
           "RegisterShuffle with ShuffleId["
               + shuffleId
@@ -158,8 +159,9 @@ public <K, V, C> ShuffleHandle registerShuffle(
 
     startHeartbeat();
 
-    shuffleIdToPartitionNum.putIfAbsent(shuffleId, dependency.partitioner().numPartitions());
-    shuffleIdToNumMapTasks.putIfAbsent(shuffleId, dependency.rdd().partitions().length);
+    shuffleIdToPartitionNum.computeIfAbsent(
+        shuffleId, key -> dependency.partitioner().numPartitions());
+    shuffleIdToNumMapTasks.computeIfAbsent(shuffleId, key -> dependency.rdd().partitions().length);
     if (shuffleManagerRpcServiceEnabled && rssStageRetryForWriteFailureEnabled) {
       ShuffleHandleInfo handleInfo =
           new MutableShuffleHandleInfo(shuffleId, partitionToServers, remoteStorage);

client-spark/spark3/src/main/java/org/apache/spark/shuffle/RssShuffleManager.java

@@ -135,7 +135,8 @@ public <K, V, C> ShuffleHandle registerShuffle(
 
     if (dependency.partitioner().numPartitions() == 0) {
       shuffleIdToPartitionNum.putIfAbsent(shuffleId, 0);
-      shuffleIdToNumMapTasks.putIfAbsent(shuffleId, dependency.rdd().partitions().length);
+      shuffleIdToNumMapTasks.computeIfAbsent(
+          shuffleId, key -> dependency.rdd().partitions().length);
       LOG.info(
           "RegisterShuffle with ShuffleId["
               + shuffleId
@@ -175,8 +176,9 @@ public <K, V, C> ShuffleHandle registerShuffle(
             rssStageResubmitManager.getServerIdBlackList(),
             0);
     startHeartbeat();
-    shuffleIdToPartitionNum.putIfAbsent(shuffleId, dependency.partitioner().numPartitions());
-    shuffleIdToNumMapTasks.putIfAbsent(shuffleId, dependency.rdd().partitions().length);
+    shuffleIdToPartitionNum.computeIfAbsent(
+        shuffleId, key -> dependency.partitioner().numPartitions());
+    shuffleIdToNumMapTasks.computeIfAbsent(shuffleId, key -> dependency.rdd().partitions().length);
     if (shuffleManagerRpcServiceEnabled && rssStageRetryForWriteFailureEnabled) {
       ShuffleHandleInfo shuffleHandleInfo =
           new MutableShuffleHandleInfo(

client-tez/src/main/java/org/apache/tez/runtime/library/common/shuffle/impl/RssShuffleManager.java

@@ -870,7 +870,7 @@ public void addKnownInput(
           pendingPartition.add(p);
         }
         allRssPartition.add(p);
-        partitionToInput.putIfAbsent(p, new ArrayList<>());
+        partitionToInput.computeIfAbsent(p, key -> new ArrayList<>());
         partitionToInput.get(p).add(srcAttemptIdentifier);
         LOG.info("Add partition:{}, after add, now partition:{}", p, allRssPartition);
       }

client-tez/src/main/java/org/apache/tez/runtime/library/common/shuffle/orderedgrouped/RssShuffleScheduler.java

@@ -1299,7 +1299,7 @@ public synchronized void addKnownMapOutput(
     partitionIdToSuccessMapTaskAttempts.get(partitionId).add(srcAttempt);
     String pathComponent = srcAttempt.getPathComponent();
     TezTaskAttemptID tezTaskAttemptId = IdUtils.convertTezTaskAttemptID(pathComponent);
-    partitionIdToSuccessTezTasks.putIfAbsent(partitionId, new HashSet<>());
+    partitionIdToSuccessTezTasks.computeIfAbsent(partitionId, key -> new HashSet<>());
     partitionIdToSuccessTezTasks.get(partitionId).add(tezTaskAttemptId.getTaskID());
 
     uniqueHosts.add(new HostPort(inputHostName, port));

coordinator/src/main/java/org/apache/uniffle/coordinator/ApplicationManager.java

@@ -276,7 +276,7 @@ public synchronized void decRemoteStorageCounter(String storagePath) {
         }
       } else {
         LOG.warn("Can't find counter for remote storage: {}", storagePath);
-        remoteStoragePathRankValue.putIfAbsent(storagePath, new RankValue(0));
+        remoteStoragePathRankValue.computeIfAbsent(storagePath, key -> new RankValue(0));
       }
       if (remoteStoragePathRankValue.get(storagePath).getAppNum().get() == 0
           && !availableRemoteStorageInfo.containsKey(storagePath)) {

server/src/main/java/org/apache/uniffle/server/merge/Shuffle.java

@@ -20,6 +20,7 @@
 import java.io.IOException;
 import java.util.Comparator;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
 
 import com.google.common.annotations.VisibleForTesting;
 import org.roaringbitmap.longlong.Roaring64NavigableMap;
@@ -69,8 +70,23 @@ public Shuffle(
 
   public void startSortMerge(int partitionId, Roaring64NavigableMap expectedBlockIdMap)
       throws IOException {
-    this.partitions.putIfAbsent(partitionId, new Partition<K, V>(this, partitionId));
-    this.partitions.get(partitionId).startSortMerge(expectedBlockIdMap);
+    AtomicReference<IOException> exception = new AtomicReference<>();
+    Partition<K, V> partition =
+        this.partitions.computeIfAbsent(
+            partitionId,
+            key -> {
+              try {
+                return new Partition<K, V>(this, partitionId);
+              } catch (IOException e) {
+                exception.set(e);
+              }
+              return null;
+            });
+    if (exception.get() != null) {
+      throw exception.get();
+    }
+    assert partition != null;
+    partition.startSortMerge(expectedBlockIdMap);
   }
 
   void cleanup() {

server/src/main/java/org/apache/uniffle/server/merge/ShuffleMergeManager.java

@@ -163,22 +163,23 @@ public StatusCode registerShuffle(String appId, int shuffleId, MergeContext merg
       } else {
         comparator = defaultComparator;
       }
-      this.shuffles.putIfAbsent(appId, JavaUtils.newConcurrentMap());
+      this.shuffles.computeIfAbsent(appId, key -> JavaUtils.newConcurrentMap());
       this.shuffles
           .get(appId)
-          .putIfAbsent(
+          .computeIfAbsent(
               shuffleId,
-              new Shuffle(
-                  serverConf,
-                  eventHandler,
-                  shuffleServer,
-                  appId,
-                  shuffleId,
-                  kClass,
-                  vClass,
-                  comparator,
-                  mergeContext.getMergedBlockSize(),
-                  classLoader));
+              key ->
+                  new Shuffle(
+                      serverConf,
+                      eventHandler,
+                      shuffleServer,
+                      appId,
+                      shuffleId,
+                      kClass,
+                      vClass,
+                      comparator,
+                      mergeContext.getMergedBlockSize(),
+                      classLoader));
     } catch (ClassNotFoundException
         | InstantiationException
         | IllegalAccessException

storage/src/main/java/org/apache/uniffle/storage/common/LocalStorageMeta.java

@@ -93,22 +93,13 @@ public void setSize(long diskSize) {
     this.size.set(diskSize);
   }
 
-  /**
-   * If the method is implemented as below:
-   *
-   * <p>if (shuffleMetaMap.contains(shuffleId)) { // `Time A` return shuffleMetaMap.get(shuffleId) }
-   * else { shuffleMetaMap.putIfAbsent(shuffleId, newMeta) return newMeta }
-   *
-   * <p>Because if shuffleMetaMap remove shuffleId at `Time A` in another thread,
-   * shuffleMetaMap.get(shuffleId) will return null. We need to guarantee that this method is thread
-   * safe, and won't return null.
-   */
   public void createMetadataIfNotExist(String shuffleKey) {
-    ShuffleMeta meta = new ShuffleMeta();
-    ShuffleMeta oldMeta = shuffleMetaMap.putIfAbsent(shuffleKey, meta);
-    if (oldMeta == null) {
-      LOG.info("Create metadata of shuffle {}.", shuffleKey);
-    }
+    shuffleMetaMap.computeIfAbsent(
+        shuffleKey,
+        key -> {
+          LOG.info("Create metadata of shuffle {}.", key);
+          return new ShuffleMeta();
+        });
   }
 
   private ShuffleMeta getShuffleMeta(String shuffleKey) {