up

Author: yew1eb

client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala

@@ -78,6 +78,9 @@ class ReducePartitionCommitHandler(
   // TODO: Move this to native Int -> Int Map
   private val shuffleToCompletedMappers = JavaUtils.newConcurrentHashMap[Int, Int]()
   private val shuffleIdLocks = JavaUtils.newConcurrentHashMap[Int, Object]()
+  private val shuffleIdLocksRegisterFunc = new util.function.Function[Int, Object] {
+    override def apply(key: Int): Object = new Object()
+  }
 
   private val stageEndTimeout = conf.clientPushStageEndTimeout
   private val mockShuffleLost = conf.testMockShuffleLost
@@ -317,7 +320,7 @@ class ReducePartitionCommitHandler(
       numPartitions: Int,
       crc32PerPartition: Array[Int],
       bytesWrittenPerPartition: Array[Long]): (Boolean, Boolean) = {
-    val shuffleLock = shuffleIdLocks.computeIfAbsent(shuffleId, _ => new Object())
+    val shuffleLock = shuffleIdLocks.computeIfAbsent(shuffleId, shuffleIdLocksRegisterFunc)
     val (mapperAttemptFinishedSuccess, allMapperFinished) = shuffleLock.synchronized {
       if (getMapperAttempts(shuffleId) == null) {
         logDebug(s"[handleMapperEnd] $shuffleId not registered, create one.")
@@ -432,7 +435,7 @@ class ReducePartitionCommitHandler(
   }
 
   private def initMapperAttempts(shuffleId: Int, numMappers: Int, numPartitions: Int): Unit = {
-    val shuffleLock = shuffleIdLocks.computeIfAbsent(shuffleId, _ => new Object())
+    val shuffleLock = shuffleIdLocks.computeIfAbsent(shuffleId, shuffleIdLocksRegisterFunc)
     shuffleLock.synchronized {
       if (!shuffleMapperAttempts.containsKey(shuffleId)) {
         val attempts = new Array[Int](numMappers)