apache · zuston · Jan 17, 2023 · Jan 11, 2023 · Jan 11, 2023 · Jan 11, 2023
diff --git a/common/src/main/java/org/apache/uniffle/common/util/TripleFunction.java b/common/src/main/java/org/apache/uniffle/common/util/TripleFunction.java
@@ -0,0 +1,25 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.uniffle.common.util;
+
+@FunctionalInterface
+public interface TripleFunction<T, U, E, R> {
+
+  R accept(T t, U u, E e);
+
+}
diff --git a/docs/server_guide.md b/docs/server_guide.md
@@ -86,6 +86,8 @@ This document will introduce how to deploy Uniffle shuffle servers.
 |rss.server.leak.shuffledata.check.interval|3600000|The interval of leak shuffle data check (ms)|
 |rss.server.max.concurrency.of.single.partition.writer|1|The max concurrency of single partition writer, the data partition file number is equal to this value. Default value is 1. This config could improve the writing speed, especially for huge partition.|
 |rss.metrics.reporter.class|-|The class of metrics reporter.|
+|rss.server.huge-partition.size.threshold|20g|Threshold of huge partition size, once exceeding threshold, memory usage limitation and huge partition buffer flushing will be triggered.|
+|rss.server.huge-partition.memory.limit.ratio|0.2|The memory usage limit ratio for huge partition, it will only triggered when partition's size exceeds the threshold of 'rss.server.huge-partition.size.threshold'|
 
 ### Advanced Configurations
 |Property Name|Default| Description                                                                                                                                                                                 |

diff --git a/...nal-client/src/main/java/org/apache/uniffle/client/impl/grpc/ShuffleServerGrpcClient.java b/...nal-client/src/main/java/org/apache/uniffle/client/impl/grpc/ShuffleServerGrpcClient.java
@@ -17,6 +17,8 @@
 
 package org.apache.uniffle.client.impl.grpc;
 
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Random;
@@ -179,8 +181,33 @@ private AppHeartBeatResponse doSendHeartBeat(String appId, long timeout) {
     return blockingStub.withDeadlineAfter(timeout, TimeUnit.MILLISECONDS).appHeartbeat(request);
   }
 
-  public long requirePreAllocation(int requireSize, int retryMax, long retryIntervalMax) {
-    RequireBufferRequest rpcRequest = RequireBufferRequest.newBuilder().setRequireSize(requireSize).build();
+  // Only for tests
+  @VisibleForTesting
+  public long requirePreAllocation(int requireSize, int retryMax, long retryIntervalMax) throws Exception {
+    return requirePreAllocation(
+        "EMPTY",
+        0,
+        Collections.emptyList(),
+        requireSize,
+        retryMax,
+        retryIntervalMax
+    );
+  }
+
+  public long requirePreAllocation(
+      String appId,
+      int shuffleId,
+      List<Integer> partitionIds,
+      int requireSize,
+      int retryMax,
+      long retryIntervalMax) {
+    RequireBufferRequest rpcRequest = RequireBufferRequest.newBuilder()
+        .setShuffleId(shuffleId)
+        .addAllPartitionIds(partitionIds)
+        .setAppId(appId)
+        .setRequireSize(requireSize)
+        .build();
+
     long start = System.currentTimeMillis();
     RequireBufferResponse rpcResponse = getBlockingStub().requireBuffer(rpcRequest);
     int retry = 0;
@@ -282,6 +309,9 @@ public RssSendShuffleDataResponse sendShuffleData(RssSendShuffleDataRequest requ
       List<ShuffleData> shuffleData = Lists.newArrayList();
       int size = 0;
       int blockNum = 0;
+      int shuffleId = stb.getKey();
+      List<Integer> partitionIds = new ArrayList<>();
+
       for (Map.Entry<Integer, List<ShuffleBlockInfo>> ptb : stb.getValue().entrySet()) {
         List<ShuffleBlock> shuffleBlocks = Lists.newArrayList();
         for (ShuffleBlockInfo sbi : ptb.getValue()) {
@@ -298,14 +328,21 @@ public RssSendShuffleDataResponse sendShuffleData(RssSendShuffleDataRequest requ
         shuffleData.add(ShuffleData.newBuilder().setPartitionId(ptb.getKey())
             .addAllBlock(shuffleBlocks)
             .build());
+        partitionIds.add(ptb.getKey());
       }
 
       final int allocateSize = size;
       final int finalBlockNum = blockNum;
       try {
         RetryUtils.retry(() -> {
-          long requireId = requirePreAllocation(allocateSize, request.getRetryMax() / maxRetryAttempts,
-              request.getRetryIntervalMax());
+          long requireId = requirePreAllocation(
+              appId,
+              shuffleId,
+              partitionIds,
+              allocateSize,
+              request.getRetryMax() / maxRetryAttempts,
+              request.getRetryIntervalMax()
+          );
           if (requireId == FAILED_REQUIRE_ID) {
             throw new RssException(String.format(
                 "requirePreAllocation failed! size[%s], host[%s], port[%s]", allocateSize, host, port));

diff --git a/proto/src/main/proto/Rss.proto b/proto/src/main/proto/Rss.proto
@@ -51,6 +51,9 @@ message FinishShuffleResponse {
 
 message RequireBufferRequest {
   int32 requireSize = 1;
+  string appId = 2;
+  int32 shuffleId = 3;
+  repeated int32 partitionIds = 4;
 }
 
 message RequireBufferResponse {

diff --git a/server/src/main/java/org/apache/uniffle/server/ShuffleServerConf.java b/server/src/main/java/org/apache/uniffle/server/ShuffleServerConf.java
@@ -342,6 +342,20 @@ public class ShuffleServerConf extends RssBaseConf {
       .noDefaultValue()
       .withDescription("The env key to get json source of local storage media provider");
 
+  public static final ConfigOption<Long> HUGE_PARTITION_SIZE_THRESHOLD = ConfigOptions
+      .key("rss.server.huge-partition.size.threshold")
+      .longType()
+      .defaultValue(20 * 1024 * 1024 * 1024L)
+      .withDescription("Threshold of huge partition size, once exceeding threshold, memory usage limitation and "
+          + "huge partition buffer flushing will be triggered.");
+
+  public static final ConfigOption<Double> HUGE_PARTITION_MEMORY_USAGE_LIMITATION_RATIO = ConfigOptions
+      .key("rss.server.huge-partition.memory.limit.ratio")
+      .doubleType()
+      .defaultValue(0.2)
+      .withDescription("The memory usage limit ratio for huge partition, it will only triggered when partition's "
+          + "size exceeds the threshold of '" + HUGE_PARTITION_SIZE_THRESHOLD.key() + "'");
+
   public ShuffleServerConf() {
   }
 

diff --git a/server/src/main/java/org/apache/uniffle/server/ShuffleServerGrpcService.java b/server/src/main/java/org/apache/uniffle/server/ShuffleServerGrpcService.java
@@ -30,6 +30,7 @@
 import io.grpc.Context;
 import io.grpc.Status;
 import io.grpc.stub.StreamObserver;
+import org.apache.commons.lang3.StringUtils;
 import org.roaringbitmap.longlong.Roaring64NavigableMap;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -353,7 +354,22 @@ public void finishShuffle(FinishShuffleRequest req,
   @Override
   public void requireBuffer(RequireBufferRequest request,
       StreamObserver<RequireBufferResponse> responseObserver) {
-    long requireBufferId = shuffleServer.getShuffleTaskManager().requireBuffer(request.getRequireSize());
+    String appId = request.getAppId();
+    long requireBufferId;
+    if (StringUtils.isEmpty(appId)) {
+      // To be compatible with older client version
+      requireBufferId = shuffleServer.getShuffleTaskManager().requireBuffer(
+          request.getRequireSize()
+      );
+    } else {
+      requireBufferId = shuffleServer.getShuffleTaskManager().requireBuffer(
+          appId,
+          request.getShuffleId(),
+          request.getPartitionIdsList(),
+          request.getRequireSize()
+      );
+    }
+
     StatusCode status = StatusCode.SUCCESS;
     if (requireBufferId == -1) {
       status = StatusCode.NO_BUFFER;

diff --git a/server/src/main/java/org/apache/uniffle/server/ShuffleTaskManager.java b/server/src/main/java/org/apache/uniffle/server/ShuffleTaskManager.java
@@ -95,7 +95,6 @@ public class ShuffleTaskManager {
   private Map<Long, PreAllocatedBufferInfo> requireBufferIds = Maps.newConcurrentMap();
   private Runnable clearResourceThread;
   private BlockingQueue<PurgeEvent> expiredAppIdQueue = Queues.newLinkedBlockingQueue();
-  // appId -> shuffleId -> serverReadHandler
 
   public ShuffleTaskManager(
       ShuffleServerConf conf,
@@ -200,7 +199,13 @@ public StatusCode registerShuffle(
   public StatusCode cacheShuffleData(
       String appId, int shuffleId, boolean isPreAllocated, ShufflePartitionedData spd) {
     refreshAppId(appId);
-    return shuffleBufferManager.cacheShuffleData(appId, shuffleId, isPreAllocated, spd);
+    return shuffleBufferManager.cacheShuffleData(
+        appId,
+        shuffleId,
+        isPreAllocated,
+        spd,
+        this::getPartitionDataSize
+    );
   }
 
   public PreAllocatedBufferInfo getAndRemovePreAllocatedBuffer(long requireBufferId) {
@@ -335,6 +340,27 @@ public Roaring64NavigableMap getCachedBlockIds(String appId, int shuffleId) {
     return blockIds;
   }
 
+  public long getPartitionDataSize(String appId, int shuffleId, int partitionId) {
+    ShuffleTaskInfo shuffleTaskInfo = shuffleTaskInfos.get(appId);
+    if (shuffleTaskInfo == null) {
+      return 0L;
+    }
+    return shuffleTaskInfo.getPartitionDataSize(shuffleId, partitionId);
+  }
+
+  public long requireBuffer(String appId, int shuffleId, List<Integer> partitionIds, int requireSize) {
+    ShuffleTaskInfo shuffleTaskInfo = shuffleTaskInfos.get(appId);
+    if (shuffleTaskInfo != null) {
+      for (int partitionId : partitionIds) {
+        long partitionUsedDataSize = getPartitionDataSize(appId, shuffleId, partitionId);
+        if (shuffleBufferManager.limitHugePartition(appId, shuffleId, partitionId, partitionUsedDataSize)) {
+          return -1;
+        }
+      }
+    }
+    return requireBuffer(requireSize);
+  }
+
   public long requireBuffer(int requireSize) {
     long requireId = -1;
     if (shuffleBufferManager.requireMemory(requireSize, true)) {
@@ -630,5 +656,4 @@ private void triggerFlush() {
       this.shuffleBufferManager.flushIfNecessary();
     }
   }
-
 }
diff --git a/server/src/main/java/org/apache/uniffle/server/buffer/ShuffleBufferManager.java b/server/src/main/java/org/apache/uniffle/server/buffer/ShuffleBufferManager.java
@@ -41,6 +41,7 @@
 import org.apache.uniffle.common.ShufflePartitionedData;
 import org.apache.uniffle.common.util.Constants;
 import org.apache.uniffle.common.util.RssUtils;
+import org.apache.uniffle.common.util.TripleFunction;
 import org.apache.uniffle.server.ShuffleDataFlushEvent;
 import org.apache.uniffle.server.ShuffleFlushManager;
 import org.apache.uniffle.server.ShuffleServerConf;
@@ -62,6 +63,9 @@ public class ShuffleBufferManager {
   // when shuffle buffer manager flushes data, shuffles with data size < shuffleFlushThreshold is kept in memory to
   // reduce small I/Os to persistent storage, especially for local HDDs.
   private long shuffleFlushThreshold;
+  // Huge partition vars
+  private long hugePartitionSizeThreshold;
+  private long hugePartitionMemoryLimitSize;
 
   protected long bufferSize = 0;
   protected AtomicLong preAllocatedSize = new AtomicLong(0L);
@@ -86,6 +90,10 @@ public ShuffleBufferManager(ShuffleServerConf conf, ShuffleFlushManager shuffleF
     this.bufferFlushEnabled = conf.getBoolean(ShuffleServerConf.SINGLE_BUFFER_FLUSH_ENABLED);
     this.bufferFlushThreshold = conf.getLong(ShuffleServerConf.SINGLE_BUFFER_FLUSH_THRESHOLD);
     this.shuffleFlushThreshold = conf.getLong(ShuffleServerConf.SERVER_SHUFFLE_FLUSH_THRESHOLD);
+    this.hugePartitionSizeThreshold = conf.getSizeAsBytes(ShuffleServerConf.HUGE_PARTITION_SIZE_THRESHOLD);
+    this.hugePartitionMemoryLimitSize = Math.round(
+        capacity * conf.get(ShuffleServerConf.HUGE_PARTITION_MEMORY_USAGE_LIMITATION_RATIO)
+    );
   }
 
   public StatusCode registerBuffer(String appId, int shuffleId, int startPartition, int endPartition) {
@@ -104,8 +112,27 @@ public StatusCode registerBuffer(String appId, int shuffleId, int startPartition
     return StatusCode.SUCCESS;
   }
 
-  public StatusCode cacheShuffleData(String appId, int shuffleId,
-      boolean isPreAllocated, ShufflePartitionedData spd) {
+  // Only for tests
+  public StatusCode cacheShuffleData(
+      String appId,
+      int shuffleId,
+      boolean isPreAllocated,
+      ShufflePartitionedData spd) {
+    return cacheShuffleData(
+        appId,
+        shuffleId,
+        isPreAllocated,
+        spd,
+        null
+    );
+  }
+
+  public StatusCode cacheShuffleData(
+      String appId,
+      int shuffleId,
+      boolean isPreAllocated,
+      ShufflePartitionedData spd,
+      TripleFunction<String, Integer, Integer, Long> getPartitionDataSizeFunc) {
     if (!isPreAllocated && isFull()) {
       LOG.warn("Got unexpected data, can't cache it because the space is full");
       return StatusCode.NO_BUFFER;
@@ -124,8 +151,15 @@ public StatusCode cacheShuffleData(String appId, int shuffleId,
     }
     updateShuffleSize(appId, shuffleId, size);
     synchronized (this) {
-      flushSingleBufferIfNecessary(buffer, appId, shuffleId,
-          entry.getKey().lowerEndpoint(), entry.getKey().upperEndpoint());
+      flushSingleBufferIfNecessary(
+          buffer,
+          appId,
+          shuffleId,
+          spd.getPartitionId(),
+          entry.getKey().lowerEndpoint(),
+          entry.getKey().upperEndpoint(),
+          getPartitionDataSizeFunc
+      );
       flushIfNecessary();
     }
     return StatusCode.SUCCESS;
@@ -184,12 +218,26 @@ public ShuffleDataResult getShuffleData(
     return buffer.getShuffleData(blockId, readBufferSize, expectedTaskIds);
   }
 
-  void flushSingleBufferIfNecessary(ShuffleBuffer buffer, String appId,
-      int shuffleId, int startPartition, int endPartition) {
+  void flushSingleBufferIfNecessary(
+      ShuffleBuffer buffer,
+      String appId,
+      int shuffleId,
+      int partitionId,
+      int startPartition,
+      int endPartition,
+      TripleFunction<String, Integer, Integer, Long> getPartitionDataSizeFunc) {
     // When we use multi storage and trigger single buffer flush, the buffer size should be bigger
     // than rss.server.flush.cold.storage.threshold.size, otherwise cold storage will be useless.
     if (this.bufferFlushEnabled && buffer.getSize() > this.bufferFlushThreshold) {
       flushBuffer(buffer, appId, shuffleId, startPartition, endPartition);
+      return;
+    }
+
+    if (getPartitionDataSizeFunc != null
+        && getPartitionDataSizeFunc.accept(appId, shuffleId, partitionId) > hugePartitionSizeThreshold
+        && buffer.getSize() > this.bufferFlushThreshold) {
+      flushBuffer(buffer, appId, shuffleId, startPartition, endPartition);
+      return;
     }
   }
 
@@ -515,4 +563,16 @@ public void removeBufferByShuffleId(String appId, Integer... shuffleIds) {
       shuffleIdToBuffers.remove(shuffleId);
     }
   }
+
+  public boolean limitHugePartition(String appId, int shuffleId, int partitionId, long usedPartitionDataSize) {
+    if (usedPartitionDataSize > hugePartitionSizeThreshold) {
+      long memoryUsed = getShuffleBufferEntry(appId, shuffleId, partitionId).getValue().getSize();
+      if (memoryUsed > hugePartitionMemoryLimitSize) {
+        LOG.warn("AppId: {}, shuffleId: {}, partitionId: {}, memory used: {}, "
+            + "huge partition triggered memory limitation.", appId, shuffleId, partitionId, memoryUsed);
+        return true;
+      }
+    }
+    return false;
+  }
 }