Fix style and add lock

apache · Jul 27, 2024 · 62c10a4 · 62c10a4
1 parent fbb9200
commit 62c10a4
Showing 1 changed file with 17 additions and 8 deletions.
diff --git a/server/src/main/java/org/apache/uniffle/server/ShuffleTaskManager.java b/server/src/main/java/org/apache/uniffle/server/ShuffleTaskManager.java
@@ -548,11 +548,18 @@ public long requireBuffer(
   public long requireBuffer(String appId, int requireSize) {
     if (shuffleBufferManager.requireMemory(requireSize, true)) {
       long requireId = requireBufferId.incrementAndGet();
-      Map<Long, PreAllocatedBufferInfo> requireBufferMaps =
-          appIdToRequireBufferIdsMap.computeIfAbsent(appId, x -> JavaUtils.newConcurrentMap());
-      requireBufferMaps.put(
-          requireId,
-          new PreAllocatedBufferInfo(appId, requireId, System.currentTimeMillis(), requireSize));
+      ReentrantReadWriteLock.WriteLock appLock = getAppWriteLock(appId);
+      try {
+        // preAllocatedBufferCheck will obtain lock and remove the empty appId
+        appLock.lock();
+        Map<Long, PreAllocatedBufferInfo> requireBufferMaps =
+            appIdToRequireBufferIdsMap.computeIfAbsent(appId, x -> JavaUtils.newConcurrentMap());
+        requireBufferMaps.put(
+            requireId,
+            new PreAllocatedBufferInfo(appId, requireId, System.currentTimeMillis(), requireSize));
+      } finally {
+        appLock.unlock();
+      }
       return requireId;
     } else {
       LOG.warn("Failed to require buffer, require size: {}", requireSize);
@@ -839,8 +846,10 @@ public void removeResources(String appId, boolean checkAppExpired) {
       shuffleBufferManager.removeBuffer(appId);
       shuffleFlushManager.removeResources(appId);
       Map<Long, PreAllocatedBufferInfo> requireBufferIdsMap = appIdToRequireBufferIdsMap.get(appId);
-      for (PreAllocatedBufferInfo info : requireBufferIdsMap.values()) {
-        removeAndReleasePreAllocatedBuffer(appId, info.getRequireId());
+      if (appIdToRequireBufferIdsMap != null) {
+        for (PreAllocatedBufferInfo info : requireBufferIdsMap.values()) {
+          removeAndReleasePreAllocatedBuffer(appId, info.getRequireId());
+        }
       }
 
       String operationMsg = String.format("removing storage data for appId:%s", appId);
@@ -913,7 +922,7 @@ private void preAllocatedBufferCheck() {
           appIdToRequireBufferIdsMap.entrySet()) {
         String appId = entry.getKey();
         if (MapUtils.isEmpty(entry.getValue())) {
-          ReentrantReadWriteLock.ReadLock appLock = getAppReadLock(appId);
+          ReentrantReadWriteLock.WriteLock appLock = getAppWriteLock(appId);
           try {
             appLock.lock();
             // After double check, remove empty map related this appId from