[SPARK-38555][NETWORK][SHUFFLE] Avoid contention and get or create cl…

…ientPools quickly in the TransportClientFactory ### What changes were proposed in this pull request? Avoid contention and get or create clientPools quickly in the TransportClientFactory. ### Why are the changes needed? Avoid contention for getting or creating clientPools, and clean up the code. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? Existing unittests. Closes apache#35860 from weixiuli/SPARK-38555-NETWORK. Authored-by: weixiuli <[email protected]> Signed-off-by: Mridul Muralidharan <mridul<at>gmail.com>
bruno0o · Mar 16, 2022 · 6d3e8eb · 6d3e8eb
1 parent 5967f29
commit 6d3e8eb
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/.../network-common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java b/.../network-common/src/main/java/org/apache/spark/network/client/TransportClientFactory.java
@@ -155,12 +155,8 @@ public TransportClient createClient(String remoteHost, int remotePort, boolean f
       InetSocketAddress.createUnresolved(remoteHost, remotePort);
 
     // Create the ClientPool if we don't have it yet.
-    ClientPool clientPool = connectionPool.get(unresolvedAddress);
-    if (clientPool == null) {
-      connectionPool.putIfAbsent(unresolvedAddress, new ClientPool(numConnectionsPerPeer));
-      clientPool = connectionPool.get(unresolvedAddress);
-    }
-
+    ClientPool clientPool = connectionPool.computeIfAbsent(unresolvedAddress,
+        key -> new ClientPool(numConnectionsPerPeer));
     int clientIndex = rand.nextInt(numConnectionsPerPeer);
     TransportClient cachedClient = clientPool.clients[clientIndex];