You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@celeborn.apache.org by zh...@apache.org on 2023/03/26 01:42:03 UTC
[incubator-celeborn] branch main updated: [CELEBORN-474] Speed up ConcurrentHashMap#computeIfAbsent (#1383)
This is an automated email from the ASF dual-hosted git repository.
zhouky pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new 7c444cb0c [CELEBORN-474] Speed up ConcurrentHashMap#computeIfAbsent (#1383)
7c444cb0c is described below
commit 7c444cb0c5c5d621d479cb7d6a9636bae3982e35
Author: Fei Wang <fw...@ebay.com>
AuthorDate: Sun Mar 26 09:41:59 2023 +0800
[CELEBORN-474] Speed up ConcurrentHashMap#computeIfAbsent (#1383)
---
.../flink/network/FlinkTransportClientFactory.java | 3 ++-
.../plugin/flink/network/ReadClientHandler.java | 5 +++--
.../flink/readclient/FlinkShuffleClientImpl.java | 4 +++-
...TransportFrameDecoderWithBufferSupplierSuiteJ.java | 3 ++-
.../celeborn/plugin/flink/RemoteShuffleMaster.java | 4 ++--
.../apache/celeborn/plugin/flink/ShuffleTaskInfo.java | 17 +++++++++++------
.../org/apache/celeborn/client/ShuffleClientImpl.java | 19 +++++++++----------
.../celeborn/client/ChangePartitionManager.scala | 4 ++--
.../apache/celeborn/client/commit/CommitHandler.scala | 4 ++--
.../client/commit/MapPartitionCommitHandler.scala | 3 ++-
.../client/commit/ReducePartitionCommitHandler.scala | 9 +++++----
.../apache/celeborn/client/DummyShuffleClient.java | 3 ++-
.../common/network/client/TransportClientFactory.java | 2 +-
.../network/client/TransportResponseHandler.java | 7 ++++---
.../common/network/server/BufferStreamManager.java | 9 +++++----
.../common/network/server/ChunkStreamManager.java | 5 +++--
.../org/apache/celeborn/common/util/JavaUtils.java | 19 +++++++++++++++++++
.../celeborn/common/write/InFlightRequestTracker.java | 3 ++-
.../org/apache/celeborn/common/write/PushState.java | 3 ++-
.../celeborn/common/write/SlowStartPushStrategy.java | 3 ++-
.../master/clustermeta/AbstractMetaManager.java | 5 +++--
.../clustermeta/ha/MasterStateMachineSuiteJ.java | 10 +++++++---
.../worker/congestcontrol/CongestionController.java | 3 ++-
.../deploy/worker/storage/PartitionFilesSorter.java | 13 +++++--------
.../deploy/worker/storage/StorageManager.scala | 4 ++--
25 files changed, 102 insertions(+), 62 deletions(-)
diff --git a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/FlinkTransportClientFactory.java b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/FlinkTransportClientFactory.java
index 58ec46562..40724ba35 100644
--- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/FlinkTransportClientFactory.java
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/FlinkTransportClientFactory.java
@@ -27,6 +27,7 @@ import org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf;
import org.apache.celeborn.common.network.TransportContext;
import org.apache.celeborn.common.network.client.TransportClient;
import org.apache.celeborn.common.network.client.TransportClientFactory;
+import org.apache.celeborn.common.util.JavaUtils;
public class FlinkTransportClientFactory extends TransportClientFactory {
@@ -34,7 +35,7 @@ public class FlinkTransportClientFactory extends TransportClientFactory {
public FlinkTransportClientFactory(TransportContext context) {
super(context);
- bufferSuppliers = new ConcurrentHashMap<>();
+ bufferSuppliers = JavaUtils.newConcurrentHashMap();
}
public TransportClient createClient(String remoteHost, int remotePort)
diff --git a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/ReadClientHandler.java b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/ReadClientHandler.java
index cdcccceff..c23bce9b1 100644
--- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/ReadClientHandler.java
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/network/ReadClientHandler.java
@@ -30,13 +30,14 @@ import org.apache.celeborn.common.network.protocol.BufferStreamEnd;
import org.apache.celeborn.common.network.protocol.RequestMessage;
import org.apache.celeborn.common.network.protocol.TransportableError;
import org.apache.celeborn.common.network.server.BaseMessageHandler;
+import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.plugin.flink.protocol.ReadData;
public class ReadClientHandler extends BaseMessageHandler {
private static Logger logger = LoggerFactory.getLogger(ReadClientHandler.class);
private ConcurrentHashMap<Long, Consumer<RequestMessage>> streamHandlers =
- new ConcurrentHashMap<>();
- private ConcurrentHashMap<Long, TransportClient> streamClients = new ConcurrentHashMap<>();
+ JavaUtils.newConcurrentHashMap();
+ private ConcurrentHashMap<Long, TransportClient> streamClients = JavaUtils.newConcurrentHashMap();
public void registerHandler(
long streamId, Consumer<RequestMessage> handle, TransportClient client) {
diff --git a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java
index 6bd86306f..321dcec15 100644
--- a/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java
+++ b/client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/readclient/FlinkShuffleClientImpl.java
@@ -50,6 +50,7 @@ import org.apache.celeborn.common.protocol.PbChangeLocationResponse;
import org.apache.celeborn.common.protocol.TransportModuleConstants;
import org.apache.celeborn.common.protocol.message.ControlMessages;
import org.apache.celeborn.common.protocol.message.StatusCode;
+import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.PbSerDeUtils;
import org.apache.celeborn.common.util.Utils;
import org.apache.celeborn.common.write.PushState;
@@ -62,7 +63,8 @@ public class FlinkShuffleClientImpl extends ShuffleClientImpl {
private static volatile boolean initialized = false;
private FlinkTransportClientFactory flinkTransportClientFactory;
private ReadClientHandler readClientHandler = new ReadClientHandler();
- private ConcurrentHashMap<String, TransportClient> currentClient = new ConcurrentHashMap<>();
+ private ConcurrentHashMap<String, TransportClient> currentClient =
+ JavaUtils.newConcurrentHashMap();
public static FlinkShuffleClientImpl get(
String driverHost, int port, CelebornConf conf, UserIdentifier userIdentifier) {
diff --git a/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/network/TransportFrameDecoderWithBufferSupplierSuiteJ.java b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/network/TransportFrameDecoderWithBufferSupplierSuiteJ.java
index 922b14d07..073972b44 100644
--- a/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/network/TransportFrameDecoderWithBufferSupplierSuiteJ.java
+++ b/client-flink/common/src/test/java/org/apache/celeborn/plugin/flink/network/TransportFrameDecoderWithBufferSupplierSuiteJ.java
@@ -34,13 +34,14 @@ import org.mockito.Mockito;
import org.apache.celeborn.common.network.protocol.BacklogAnnouncement;
import org.apache.celeborn.common.network.protocol.Message;
import org.apache.celeborn.common.network.protocol.ReadData;
+import org.apache.celeborn.common.util.JavaUtils;
public class TransportFrameDecoderWithBufferSupplierSuiteJ {
@Test
public void testDropUnusedBytes() throws IOException {
ConcurrentHashMap<Long, Supplier<org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf>>
- supplier = new ConcurrentHashMap<>();
+ supplier = JavaUtils.newConcurrentHashMap();
List<org.apache.flink.shaded.netty4.io.netty.buffer.ByteBuf> buffers = new ArrayList<>();
supplier.put(
diff --git a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleMaster.java b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleMaster.java
index 0d220bed8..ec9827edf 100644
--- a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleMaster.java
+++ b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleMaster.java
@@ -21,7 +21,6 @@ import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import org.apache.flink.api.common.JobID;
@@ -36,6 +35,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.client.LifecycleManager;
import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.plugin.flink.utils.FlinkUtils;
import org.apache.celeborn.plugin.flink.utils.ThreadUtils;
@@ -43,7 +43,7 @@ public class RemoteShuffleMaster implements ShuffleMaster<RemoteShuffleDescripto
private static final Logger LOG = LoggerFactory.getLogger(RemoteShuffleMaster.class);
private final ShuffleMasterContext shuffleMasterContext;
// Flink JobId -> Celeborn register shuffleIds
- private Map<JobID, Set<Integer>> jobShuffleIds = new ConcurrentHashMap<>();
+ private Map<JobID, Set<Integer>> jobShuffleIds = JavaUtils.newConcurrentHashMap();
private String celebornAppId;
private volatile LifecycleManager lifecycleManager;
private ShuffleTaskInfo shuffleTaskInfo = new ShuffleTaskInfo();
diff --git a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/ShuffleTaskInfo.java b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/ShuffleTaskInfo.java
index 4c4c2a9d1..5106cc7ae 100644
--- a/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/ShuffleTaskInfo.java
+++ b/client-flink/flink-1.14/src/main/java/org/apache/celeborn/plugin/flink/ShuffleTaskInfo.java
@@ -19,18 +19,22 @@ package org.apache.celeborn.plugin.flink;
import java.util.concurrent.ConcurrentHashMap;
+import org.apache.celeborn.common.util.JavaUtils;
+
public class ShuffleTaskInfo {
private int currentShuffleIndex = 0;
// task shuffle id -> mapId_taskAttemptId -> attemptIdx
private ConcurrentHashMap<String, ConcurrentHashMap<String, Integer>>
- taskShuffleAttemptIdToAttemptId = new ConcurrentHashMap<>();
+ taskShuffleAttemptIdToAttemptId = JavaUtils.newConcurrentHashMap();
// map attemptId index
private ConcurrentHashMap<String, ConcurrentHashMap<Integer, Integer>> taskShuffleAttemptIdIndex =
- new ConcurrentHashMap<>();
+ JavaUtils.newConcurrentHashMap();
// task shuffle id -> celeborn shuffle id
- private ConcurrentHashMap<String, Integer> taskShuffleIdToShuffleId = new ConcurrentHashMap<>();
+ private ConcurrentHashMap<String, Integer> taskShuffleIdToShuffleId =
+ JavaUtils.newConcurrentHashMap();
// celeborn shuffle id -> task shuffle id
- private ConcurrentHashMap<Integer, String> shuffleIdToTaskShuffleId = new ConcurrentHashMap<>();
+ private ConcurrentHashMap<Integer, String> shuffleIdToTaskShuffleId =
+ JavaUtils.newConcurrentHashMap();
public int getShuffleId(String taskShuffleId) {
synchronized (taskShuffleIdToShuffleId) {
@@ -48,10 +52,11 @@ public class ShuffleTaskInfo {
public int getAttemptId(String taskShuffleId, int mapId, String attemptId) {
ConcurrentHashMap<Integer, Integer> attemptIndex =
- taskShuffleAttemptIdIndex.computeIfAbsent(taskShuffleId, (id) -> new ConcurrentHashMap<>());
+ taskShuffleAttemptIdIndex.computeIfAbsent(
+ taskShuffleId, (id) -> JavaUtils.newConcurrentHashMap());
ConcurrentHashMap<String, Integer> attemptIdMap =
taskShuffleAttemptIdToAttemptId.computeIfAbsent(
- taskShuffleId, (id) -> new ConcurrentHashMap<>());
+ taskShuffleId, (id) -> JavaUtils.newConcurrentHashMap());
String mapAttemptId = mapId + "_" + attemptId;
synchronized (attemptIndex) {
if (!attemptIdMap.containsKey(mapAttemptId)) {
diff --git a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
index efb3249c4..5f84d6921 100644
--- a/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
+++ b/client/src/main/java/org/apache/celeborn/client/ShuffleClientImpl.java
@@ -55,10 +55,7 @@ import org.apache.celeborn.common.rpc.RpcAddress;
import org.apache.celeborn.common.rpc.RpcEndpointRef;
import org.apache.celeborn.common.rpc.RpcEnv;
import org.apache.celeborn.common.unsafe.Platform;
-import org.apache.celeborn.common.util.PackedPartitionId;
-import org.apache.celeborn.common.util.PbSerDeUtils;
-import org.apache.celeborn.common.util.ThreadUtils;
-import org.apache.celeborn.common.util.Utils;
+import org.apache.celeborn.common.util.*;
import org.apache.celeborn.common.write.DataBatches;
import org.apache.celeborn.common.write.PushState;
@@ -90,17 +87,18 @@ public class ShuffleClientImpl extends ShuffleClient {
// key: shuffleId, value: (partitionId, PartitionLocation)
private final Map<Integer, ConcurrentHashMap<Integer, PartitionLocation>> reducePartitionMap =
- new ConcurrentHashMap<>();
+ JavaUtils.newConcurrentHashMap();
- protected final ConcurrentHashMap<Integer, Set<String>> mapperEndMap = new ConcurrentHashMap<>();
+ protected final ConcurrentHashMap<Integer, Set<String>> mapperEndMap =
+ JavaUtils.newConcurrentHashMap();
// key: shuffleId-mapId-attemptId
- protected final Map<String, PushState> pushStates = new ConcurrentHashMap<>();
+ protected final Map<String, PushState> pushStates = JavaUtils.newConcurrentHashMap();
private final ExecutorService pushDataRetryPool;
private final ExecutorService partitionSplitPool;
- private final Map<Integer, Set<Integer>> splitting = new ConcurrentHashMap<>();
+ private final Map<Integer, Set<Integer>> splitting = JavaUtils.newConcurrentHashMap();
ThreadLocal<Compressor> compressorThreadLocal =
new ThreadLocal<Compressor>() {
@@ -138,7 +136,8 @@ public class ShuffleClientImpl extends ShuffleClient {
}
// key: shuffleId
- protected final Map<Integer, ReduceFileGroups> reduceFileGroupsMap = new ConcurrentHashMap<>();
+ protected final Map<Integer, ReduceFileGroups> reduceFileGroupsMap =
+ JavaUtils.newConcurrentHashMap();
public ShuffleClientImpl(CelebornConf conf, UserIdentifier userIdentifier) {
super();
@@ -413,7 +412,7 @@ public class ShuffleClientImpl extends ShuffleClient {
PbRegisterShuffleResponse response = callable.call();
StatusCode respStatus = Utils.toStatusCode(response.getStatus());
if (StatusCode.SUCCESS.equals(respStatus)) {
- ConcurrentHashMap<Integer, PartitionLocation> result = new ConcurrentHashMap<>();
+ ConcurrentHashMap<Integer, PartitionLocation> result = JavaUtils.newConcurrentHashMap();
for (int i = 0; i < response.getPartitionLocationsList().size(); i++) {
PartitionLocation partitionLoc =
PbSerDeUtils.fromPbPartitionLocation(response.getPartitionLocationsList().get(i));
diff --git a/client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala b/client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala
index 17b74cf97..3fcabb58b 100644
--- a/client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala
@@ -30,7 +30,7 @@ import org.apache.celeborn.common.meta.WorkerInfo
import org.apache.celeborn.common.protocol.PartitionLocation
import org.apache.celeborn.common.protocol.message.ControlMessages.WorkerResource
import org.apache.celeborn.common.protocol.message.StatusCode
-import org.apache.celeborn.common.util.{ThreadUtils, Utils}
+import org.apache.celeborn.common.util.{JavaUtils, ThreadUtils, Utils}
case class ChangePartitionRequest(
context: RequestLocationCallContext,
@@ -121,7 +121,7 @@ class ChangePartitionManager(
Int,
ConcurrentHashMap[Integer, util.Set[ChangePartitionRequest]]]() {
override def apply(s: Int): ConcurrentHashMap[Integer, util.Set[ChangePartitionRequest]] =
- new ConcurrentHashMap()
+ JavaUtils.newConcurrentHashMap()
}
private val inBatchShuffleIdRegisterFunc = new util.function.Function[Int, util.Set[Integer]]() {
diff --git a/client/src/main/scala/org/apache/celeborn/client/commit/CommitHandler.scala b/client/src/main/scala/org/apache/celeborn/client/commit/CommitHandler.scala
index 7a788343e..89632bd7d 100644
--- a/client/src/main/scala/org/apache/celeborn/client/commit/CommitHandler.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/commit/CommitHandler.scala
@@ -34,7 +34,7 @@ import org.apache.celeborn.common.protocol.{PartitionLocation, PartitionType}
import org.apache.celeborn.common.protocol.message.ControlMessages.{CommitFiles, CommitFilesResponse}
import org.apache.celeborn.common.protocol.message.StatusCode
import org.apache.celeborn.common.rpc.{RpcCallContext, RpcEndpointRef}
-import org.apache.celeborn.common.util.{CollectionUtils, ThreadUtils, Utils}
+import org.apache.celeborn.common.util.{CollectionUtils, JavaUtils, ThreadUtils, Utils}
// Can Remove this if celeborn don't support scala211 in future
import org.apache.celeborn.common.util.FunctionConverter._
@@ -176,7 +176,7 @@ abstract class CommitHandler(
recordWorkerFailure: ShuffleFailedWorkers => Unit): (Boolean, Boolean)
def registerShuffle(shuffleId: Int, numMappers: Int): Unit = {
- reducerFileGroupsMap.put(shuffleId, new ConcurrentHashMap())
+ reducerFileGroupsMap.put(shuffleId, JavaUtils.newConcurrentHashMap())
}
def parallelCommitFiles(
diff --git a/client/src/main/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandler.scala b/client/src/main/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandler.scala
index 8fbfc477a..890559981 100644
--- a/client/src/main/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandler.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/commit/MapPartitionCommitHandler.scala
@@ -37,6 +37,7 @@ import org.apache.celeborn.common.protocol.message.StatusCode
import org.apache.celeborn.common.rpc.RpcCallContext
// Can Remove this if celeborn don't support scala211 in future
import org.apache.celeborn.common.util.FunctionConverter._
+import org.apache.celeborn.common.util.JavaUtils
import org.apache.celeborn.common.util.Utils
/**
@@ -217,7 +218,7 @@ class MapPartitionCommitHandler(
context.reply(GetReducerFileGroupResponse(
StatusCode.SUCCESS,
- reducerFileGroupsMap.getOrDefault(shuffleId, new ConcurrentHashMap()),
+ reducerFileGroupsMap.getOrDefault(shuffleId, JavaUtils.newConcurrentHashMap()),
getMapperAttempts(shuffleId),
succeedPartitionIds))
}
diff --git a/client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala b/client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala
index 97929a832..b8509e971 100644
--- a/client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala
+++ b/client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala
@@ -37,6 +37,7 @@ import org.apache.celeborn.common.protocol.message.ControlMessages.GetReducerFil
import org.apache.celeborn.common.protocol.message.StatusCode
import org.apache.celeborn.common.rpc.RpcCallContext
import org.apache.celeborn.common.rpc.netty.{LocalNettyRpcCallContext, RemoteNettyRpcCallContext}
+import org.apache.celeborn.common.util.JavaUtils
/**
* This commit handler is for ReducePartition ShuffleType, which means that a Reduce Partition contains all data
@@ -254,7 +255,7 @@ class ReducePartitionCommitHandler(
logError(s"[handleGetReducerFileGroup] Wait for handleStageEnd Timeout! $shuffleId.")
context.reply(GetReducerFileGroupResponse(
StatusCode.STAGE_END_TIME_OUT,
- new ConcurrentHashMap(),
+ JavaUtils.newConcurrentHashMap(),
Array.empty))
} else {
logDebug("[handleGetReducerFileGroup] Wait for handleStageEnd complete cost" +
@@ -263,14 +264,14 @@ class ReducePartitionCommitHandler(
context.reply(
GetReducerFileGroupResponse(
StatusCode.SHUFFLE_DATA_LOST,
- new ConcurrentHashMap(),
+ JavaUtils.newConcurrentHashMap(),
Array.empty))
} else {
// LocalNettyRpcCallContext is for the UTs
if (context.isInstanceOf[LocalNettyRpcCallContext]) {
context.reply(GetReducerFileGroupResponse(
StatusCode.SUCCESS,
- reducerFileGroupsMap.getOrDefault(shuffleId, new ConcurrentHashMap()),
+ reducerFileGroupsMap.getOrDefault(shuffleId, JavaUtils.newConcurrentHashMap()),
getMapperAttempts(shuffleId)))
} else {
val cachedMsg = getReducerFileGroupRpcCache.get(
@@ -279,7 +280,7 @@ class ReducePartitionCommitHandler(
override def call(): ByteBuffer = {
val returnedMsg = GetReducerFileGroupResponse(
StatusCode.SUCCESS,
- reducerFileGroupsMap.getOrDefault(shuffleId, new ConcurrentHashMap()),
+ reducerFileGroupsMap.getOrDefault(shuffleId, JavaUtils.newConcurrentHashMap()),
getMapperAttempts(shuffleId))
context.asInstanceOf[RemoteNettyRpcCallContext].nettyEnv.serialize(returnedMsg)
}
diff --git a/client/src/test/java/org/apache/celeborn/client/DummyShuffleClient.java b/client/src/test/java/org/apache/celeborn/client/DummyShuffleClient.java
index 088dc5f21..9c88fd683 100644
--- a/client/src/test/java/org/apache/celeborn/client/DummyShuffleClient.java
+++ b/client/src/test/java/org/apache/celeborn/client/DummyShuffleClient.java
@@ -37,6 +37,7 @@ import org.apache.celeborn.client.read.RssInputStream;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.protocol.PartitionLocation;
import org.apache.celeborn.common.rpc.RpcEndpointRef;
+import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.write.PushState;
public class DummyShuffleClient extends ShuffleClient {
@@ -166,7 +167,7 @@ public class DummyShuffleClient extends ShuffleClient {
}
public void initReducePartitionMap(int shuffleId, int numPartitions, int workerNum) {
- ConcurrentHashMap<Integer, PartitionLocation> map = new ConcurrentHashMap<>();
+ ConcurrentHashMap<Integer, PartitionLocation> map = JavaUtils.newConcurrentHashMap();
String host = "host";
List<PartitionLocation> partitionLocationList = new ArrayList<>();
for (int i = 0; i < workerNum; i++) {
diff --git a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java
index 79711111d..fece7cb67 100644
--- a/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/client/TransportClientFactory.java
@@ -81,7 +81,7 @@ public class TransportClientFactory implements Closeable {
public TransportClientFactory(TransportContext context) {
this.context = Preconditions.checkNotNull(context);
this.conf = context.getConf();
- this.connectionPool = new ConcurrentHashMap<>();
+ this.connectionPool = JavaUtils.newConcurrentHashMap();
this.numConnectionsPerPeer = conf.numConnectionsPerPeer();
this.rand = new Random();
diff --git a/common/src/main/java/org/apache/celeborn/common/network/client/TransportResponseHandler.java b/common/src/main/java/org/apache/celeborn/common/network/client/TransportResponseHandler.java
index fd19f1928..46f487a2c 100644
--- a/common/src/main/java/org/apache/celeborn/common/network/client/TransportResponseHandler.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/client/TransportResponseHandler.java
@@ -37,6 +37,7 @@ import org.apache.celeborn.common.network.util.NettyUtils;
import org.apache.celeborn.common.network.util.TransportConf;
import org.apache.celeborn.common.protocol.TransportModuleConstants;
import org.apache.celeborn.common.protocol.message.StatusCode;
+import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.ThreadUtils;
import org.apache.celeborn.common.write.PushRequestInfo;
@@ -67,9 +68,9 @@ public class TransportResponseHandler extends MessageHandler<ResponseMessage> {
public TransportResponseHandler(TransportConf conf, Channel channel) {
this.conf = conf;
this.channel = channel;
- this.outstandingFetches = new ConcurrentHashMap<>();
- this.outstandingRpcs = new ConcurrentHashMap<>();
- this.outstandingPushes = new ConcurrentHashMap<>();
+ this.outstandingFetches = JavaUtils.newConcurrentHashMap();
+ this.outstandingRpcs = JavaUtils.newConcurrentHashMap();
+ this.outstandingPushes = JavaUtils.newConcurrentHashMap();
this.timeOfLastRequestNs = new AtomicLong(0);
pushTimeoutCheckerInterval = conf.pushDataTimeoutCheckIntervalMs();
synchronized (TransportResponseHandler.class) {
diff --git a/common/src/main/java/org/apache/celeborn/common/network/server/BufferStreamManager.java b/common/src/main/java/org/apache/celeborn/common/network/server/BufferStreamManager.java
index 6bbf8ab15..a4558c570 100644
--- a/common/src/main/java/org/apache/celeborn/common/network/server/BufferStreamManager.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/server/BufferStreamManager.java
@@ -54,6 +54,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.meta.FileInfo;
import org.apache.celeborn.common.network.server.memory.BufferRecycler;
import org.apache.celeborn.common.network.server.memory.MemoryManager;
+import org.apache.celeborn.common.util.JavaUtils;
public class BufferStreamManager {
private static final Logger logger = LoggerFactory.getLogger(BufferStreamManager.class);
@@ -93,9 +94,9 @@ public class BufferStreamManager {
public BufferStreamManager(int minReadBuffers, int maxReadBuffers, int threadsPerMountpoint) {
nextStreamId = new AtomicLong((long) new Random().nextInt(Integer.MAX_VALUE) * 1000);
- streams = new ConcurrentHashMap<>();
- servingStreams = new ConcurrentHashMap<>();
- activeMapPartitions = new ConcurrentHashMap<>();
+ streams = JavaUtils.newConcurrentHashMap();
+ servingStreams = JavaUtils.newConcurrentHashMap();
+ activeMapPartitions = JavaUtils.newConcurrentHashMap();
this.minReadBuffers = minReadBuffers;
this.maxReadBuffers = maxReadBuffers;
this.threadsPerMountPoint = threadsPerMountpoint;
@@ -285,7 +286,7 @@ public class BufferStreamManager {
private final Set<DataPartitionReader> readers = new HashSet<>();
private final ExecutorService readExecutor;
private final ConcurrentHashMap<Long, DataPartitionReader> streamReaders =
- new ConcurrentHashMap<>();
+ JavaUtils.newConcurrentHashMap();
/** All available buffers can be used by the partition readers for reading. */
private Queue<ByteBuf> buffers;
diff --git a/common/src/main/java/org/apache/celeborn/common/network/server/ChunkStreamManager.java b/common/src/main/java/org/apache/celeborn/common/network/server/ChunkStreamManager.java
index c9bbf5a79..83a3422fe 100644
--- a/common/src/main/java/org/apache/celeborn/common/network/server/ChunkStreamManager.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/server/ChunkStreamManager.java
@@ -32,6 +32,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.meta.FileManagedBuffers;
import org.apache.celeborn.common.meta.TimeWindow;
import org.apache.celeborn.common.network.buffer.ManagedBuffer;
+import org.apache.celeborn.common.util.JavaUtils;
/**
* StreamManager which allows registration of an Iterator<ManagedBuffer>, which are
@@ -66,8 +67,8 @@ public class ChunkStreamManager {
// For debugging purposes, start with a random stream id to help identifying different streams.
// This does not need to be globally unique, only unique to this class.
nextStreamId = new AtomicLong((long) new Random().nextInt(Integer.MAX_VALUE) * 1000);
- streams = new ConcurrentHashMap<>();
- shuffleStreamIds = new ConcurrentHashMap<>();
+ streams = JavaUtils.newConcurrentHashMap();
+ shuffleStreamIds = JavaUtils.newConcurrentHashMap();
}
public ManagedBuffer getChunk(long streamId, int chunkIndex, int offset, int len) {
diff --git a/common/src/main/java/org/apache/celeborn/common/util/JavaUtils.java b/common/src/main/java/org/apache/celeborn/common/util/JavaUtils.java
index ef284e3a6..f85c7fbd8 100644
--- a/common/src/main/java/org/apache/celeborn/common/util/JavaUtils.java
+++ b/common/src/main/java/org/apache/celeborn/common/util/JavaUtils.java
@@ -23,7 +23,9 @@ import java.nio.channels.ReadableByteChannel;
import java.nio.charset.StandardCharsets;
import java.util.Locale;
import java.util.Random;
+import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
@@ -418,4 +420,21 @@ public class JavaUtils {
}
}
}
+
+ /**
+ * For jdk8, there is bug for ConcurrentHashMap#computeIfAbsent, checking the key existence to
+ * speed up. See details in CELEBORN-474.
+ */
+ public static <K, V> ConcurrentHashMap<K, V> newConcurrentHashMap() {
+ return new ConcurrentHashMap<K, V>() {
+ @Override
+ public V computeIfAbsent(K key, Function<? super K, ? extends V> mappingFunction) {
+ V result;
+ if (null == (result = get(key))) {
+ result = super.computeIfAbsent(key, mappingFunction);
+ }
+ return result;
+ }
+ };
+ }
}
diff --git a/common/src/main/java/org/apache/celeborn/common/write/InFlightRequestTracker.java b/common/src/main/java/org/apache/celeborn/common/write/InFlightRequestTracker.java
index c311d9107..aec6bd7bf 100644
--- a/common/src/main/java/org/apache/celeborn/common/write/InFlightRequestTracker.java
+++ b/common/src/main/java/org/apache/celeborn/common/write/InFlightRequestTracker.java
@@ -28,6 +28,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.exception.CelebornIOException;
+import org.apache.celeborn.common.util.JavaUtils;
/*
* This class is for track in flight request and limit request.
@@ -42,7 +43,7 @@ public class InFlightRequestTracker {
private final AtomicInteger batchId = new AtomicInteger();
private final ConcurrentHashMap<String, Set<Integer>> inflightBatchesPerAddress =
- new ConcurrentHashMap<>();
+ JavaUtils.newConcurrentHashMap();
public InFlightRequestTracker(CelebornConf conf, PushState pushState) {
this.waitInflightTimeoutMs = conf.pushLimitInFlightTimeoutMs();
diff --git a/common/src/main/java/org/apache/celeborn/common/write/PushState.java b/common/src/main/java/org/apache/celeborn/common/write/PushState.java
index 550989ca3..e8fb9a2e9 100644
--- a/common/src/main/java/org/apache/celeborn/common/write/PushState.java
+++ b/common/src/main/java/org/apache/celeborn/common/write/PushState.java
@@ -23,6 +23,7 @@ import java.util.concurrent.atomic.AtomicReference;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.util.JavaUtils;
public class PushState {
@@ -40,7 +41,7 @@ public class PushState {
}
// key: ${master addr}-${slave addr} value: list of data batch
- public final ConcurrentHashMap<String, DataBatches> batchesMap = new ConcurrentHashMap<>();
+ public final ConcurrentHashMap<String, DataBatches> batchesMap = JavaUtils.newConcurrentHashMap();
/**
* Not thread-safe
diff --git a/common/src/main/java/org/apache/celeborn/common/write/SlowStartPushStrategy.java b/common/src/main/java/org/apache/celeborn/common/write/SlowStartPushStrategy.java
index 2517507d3..ab43cd01a 100644
--- a/common/src/main/java/org/apache/celeborn/common/write/SlowStartPushStrategy.java
+++ b/common/src/main/java/org/apache/celeborn/common/write/SlowStartPushStrategy.java
@@ -27,6 +27,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.exception.CelebornIOException;
+import org.apache.celeborn.common.util.JavaUtils;
public class SlowStartPushStrategy extends PushStrategy {
@@ -92,7 +93,7 @@ public class SlowStartPushStrategy extends PushStrategy {
this.maxInFlight = conf.pushMaxReqsInFlight();
this.initialSleepMills = conf.pushSlowStartInitialSleepTime();
this.maxSleepMills = conf.pushSlowStartMaxSleepMills();
- this.congestControlInfoPerAddress = new ConcurrentHashMap<>();
+ this.congestControlInfoPerAddress = JavaUtils.newConcurrentHashMap();
}
@VisibleForTesting
diff --git a/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/AbstractMetaManager.java b/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/AbstractMetaManager.java
index f9cff9235..596fffc4b 100644
--- a/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/AbstractMetaManager.java
+++ b/master/src/main/java/org/apache/celeborn/service/deploy/master/clustermeta/AbstractMetaManager.java
@@ -42,6 +42,7 @@ import org.apache.celeborn.common.protocol.PbSnapshotMetaInfo;
import org.apache.celeborn.common.quota.ResourceConsumption;
import org.apache.celeborn.common.rpc.RpcAddress;
import org.apache.celeborn.common.rpc.RpcEnv;
+import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.PbSerDeUtils;
import org.apache.celeborn.common.util.Utils;
@@ -52,8 +53,8 @@ public abstract class AbstractMetaManager implements IMetadataHandler {
public final Set<String> registeredShuffle = ConcurrentHashMap.newKeySet();
public final Set<String> hostnameSet = ConcurrentHashMap.newKeySet();
public final ArrayList<WorkerInfo> workers = new ArrayList<>();
- public final ConcurrentHashMap<WorkerInfo, Long> lostWorkers = new ConcurrentHashMap<>();
- public final ConcurrentHashMap<String, Long> appHeartbeatTime = new ConcurrentHashMap<>();
+ public final ConcurrentHashMap<WorkerInfo, Long> lostWorkers = JavaUtils.newConcurrentHashMap();
+ public final ConcurrentHashMap<String, Long> appHeartbeatTime = JavaUtils.newConcurrentHashMap();
// blacklist
public final Set<WorkerInfo> blacklist = ConcurrentHashMap.newKeySet();
// workerLost events
diff --git a/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/MasterStateMachineSuiteJ.java b/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/MasterStateMachineSuiteJ.java
index 7c090c55e..104ed2a0f 100644
--- a/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/MasterStateMachineSuiteJ.java
+++ b/master/src/test/java/org/apache/celeborn/service/deploy/master/clustermeta/ha/MasterStateMachineSuiteJ.java
@@ -43,6 +43,7 @@ import org.apache.celeborn.common.meta.AppDiskUsageSnapShot;
import org.apache.celeborn.common.meta.DiskInfo;
import org.apache.celeborn.common.meta.WorkerInfo;
import org.apache.celeborn.common.quota.ResourceConsumption;
+import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.Utils;
import org.apache.celeborn.service.deploy.master.clustermeta.ResourceProtos;
import org.apache.celeborn.service.deploy.master.clustermeta.ResourceProtos.RequestSlotsRequest;
@@ -210,7 +211,8 @@ public class MasterStateMachineSuiteJ extends RatisBaseSuiteJ {
disks1.put("disk1", new DiskInfo("disk1", 64 * 1024 * 1024 * 1024, 100, 100, 0));
disks1.put("disk2", new DiskInfo("disk2", 64 * 1024 * 1024 * 1024, 100, 100, 0));
disks1.put("disk3", new DiskInfo("disk3", 64 * 1024 * 1024 * 1024, 100, 100, 0));
- Map<UserIdentifier, ResourceConsumption> userResourceConsumption1 = new ConcurrentHashMap<>();
+ Map<UserIdentifier, ResourceConsumption> userResourceConsumption1 =
+ JavaUtils.newConcurrentHashMap();
userResourceConsumption1.put(
new UserIdentifier("tenant1", "name1"), new ResourceConsumption(1000, 1, 1000, 1));
userResourceConsumption1.put(
@@ -222,7 +224,8 @@ public class MasterStateMachineSuiteJ extends RatisBaseSuiteJ {
disks2.put("disk1", new DiskInfo("disk1", 64 * 1024 * 1024 * 1024, 100, 100, 0));
disks2.put("disk2", new DiskInfo("disk2", 64 * 1024 * 1024 * 1024, 100, 100, 0));
disks2.put("disk3", new DiskInfo("disk3", 64 * 1024 * 1024 * 1024, 100, 100, 0));
- Map<UserIdentifier, ResourceConsumption> userResourceConsumption2 = new ConcurrentHashMap<>();
+ Map<UserIdentifier, ResourceConsumption> userResourceConsumption2 =
+ JavaUtils.newConcurrentHashMap();
userResourceConsumption2.put(
new UserIdentifier("tenant2", "name1"), new ResourceConsumption(1000, 1, 1000, 1));
userResourceConsumption2.put(
@@ -234,7 +237,8 @@ public class MasterStateMachineSuiteJ extends RatisBaseSuiteJ {
disks3.put("disk1", new DiskInfo("disk1", 64 * 1024 * 1024 * 1024, 100, 100, 0));
disks3.put("disk2", new DiskInfo("disk2", 64 * 1024 * 1024 * 1024, 100, 100, 0));
disks3.put("disk3", new DiskInfo("disk3", 64 * 1024 * 1024 * 1024, 100, 100, 0));
- Map<UserIdentifier, ResourceConsumption> userResourceConsumption3 = new ConcurrentHashMap<>();
+ Map<UserIdentifier, ResourceConsumption> userResourceConsumption3 =
+ JavaUtils.newConcurrentHashMap();
userResourceConsumption3.put(
new UserIdentifier("tenant3", "name1"), new ResourceConsumption(1000, 1, 1000, 1));
userResourceConsumption3.put(
diff --git a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/congestcontrol/CongestionController.java b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/congestcontrol/CongestionController.java
index f6a13ef13..a7d8b0705 100644
--- a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/congestcontrol/CongestionController.java
+++ b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/congestcontrol/CongestionController.java
@@ -29,6 +29,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.identity.UserIdentifier;
import org.apache.celeborn.common.network.server.memory.MemoryManager;
+import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.ThreadUtils;
import org.apache.celeborn.service.deploy.worker.WorkerSource;
@@ -65,7 +66,7 @@ public class CongestionController {
this.lowWatermark = lowWatermark;
this.userInactiveTimeMills = userInactiveTimeMills;
this.consumedBufferStatusHub = new BufferStatusHub(sampleTimeWindowSeconds);
- this.userBufferStatuses = new ConcurrentHashMap<>();
+ this.userBufferStatuses = JavaUtils.newConcurrentHashMap();
this.removeUserExecutorService =
ThreadUtils.newDaemonSingleThreadScheduledExecutor("remove-inactive-user");
diff --git a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
index cb2670e54..485fc801e 100644
--- a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
+++ b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
@@ -54,11 +54,8 @@ import org.apache.celeborn.common.meta.FileInfo;
import org.apache.celeborn.common.metrics.source.AbstractSource;
import org.apache.celeborn.common.network.server.memory.MemoryManager;
import org.apache.celeborn.common.unsafe.Platform;
-import org.apache.celeborn.common.util.PbSerDeUtils;
-import org.apache.celeborn.common.util.ShuffleBlockInfoUtils;
+import org.apache.celeborn.common.util.*;
import org.apache.celeborn.common.util.ShuffleBlockInfoUtils.ShuffleBlockInfo;
-import org.apache.celeborn.common.util.ThreadUtils;
-import org.apache.celeborn.common.util.Utils;
import org.apache.celeborn.service.deploy.worker.LevelDBProvider;
import org.apache.celeborn.service.deploy.worker.ShuffleRecoverHelper;
import org.apache.celeborn.service.deploy.worker.WorkerSource;
@@ -72,11 +69,11 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
private File recoverFile;
private volatile boolean shutdown = false;
private final ConcurrentHashMap<String, Set<String>> sortedShuffleFiles =
- new ConcurrentHashMap<>();
+ JavaUtils.newConcurrentHashMap();
private final ConcurrentHashMap<String, Set<String>> sortingShuffleFiles =
- new ConcurrentHashMap<>();
+ JavaUtils.newConcurrentHashMap();
private final ConcurrentHashMap<String, Map<String, Map<Integer, List<ShuffleBlockInfo>>>>
- cachedIndexMaps = new ConcurrentHashMap<>();
+ cachedIndexMaps = JavaUtils.newConcurrentHashMap();
private final LinkedBlockingQueue<FileSorter> shuffleSortTaskDeque = new LinkedBlockingQueue<>();
private final AtomicInteger sortedFileCount = new AtomicInteger();
@@ -470,7 +467,7 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
indexBuf.rewind();
indexMap = ShuffleBlockInfoUtils.parseShuffleBlockInfosFromByteBuffer(indexBuf);
Map<String, Map<Integer, List<ShuffleBlockInfo>>> cacheMap =
- cachedIndexMaps.computeIfAbsent(shuffleKey, v -> new ConcurrentHashMap<>());
+ cachedIndexMaps.computeIfAbsent(shuffleKey, v -> JavaUtils.newConcurrentHashMap());
cacheMap.put(fileId, indexMap);
} catch (Exception e) {
logger.error("Read sorted shuffle file index " + indexFilePath + " error, detail: ", e);
diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
index 32c097ca4..5febb50cb 100644
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
@@ -43,7 +43,7 @@ import org.apache.celeborn.common.metrics.source.AbstractSource
import org.apache.celeborn.common.network.server.memory.MemoryManager.MemoryPressureListener
import org.apache.celeborn.common.protocol.{PartitionLocation, PartitionSplitMode, PartitionType}
import org.apache.celeborn.common.quota.ResourceConsumption
-import org.apache.celeborn.common.util.{PbSerDeUtils, ThreadUtils, Utils}
+import org.apache.celeborn.common.util.{JavaUtils, PbSerDeUtils, ThreadUtils, Utils}
import org.apache.celeborn.service.deploy.worker._
import org.apache.celeborn.service.deploy.worker.storage.StorageManager.hadoopFs
@@ -248,7 +248,7 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
private val newMapFunc =
new java.util.function.Function[String, ConcurrentHashMap[String, FileInfo]]() {
override def apply(key: String): ConcurrentHashMap[String, FileInfo] =
- new ConcurrentHashMap()
+ JavaUtils.newConcurrentHashMap()
}
private val workingDirWriterListFunc =