fix: strictly verify whether app is registered especially after serve…

…r restart
zuston · Aug 8, 2023 · bc11e2c · bc11e2c
1 parent ca24531
commit bc11e2c
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 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
@@ -409,10 +409,15 @@ public long getPartitionDataSize(String appId, int shuffleId, int partitionId) {
 
  public long requireBuffer(
  String appId, int shuffleId, List<Integer> partitionIds, int requireSize) {
- ShuffleTaskInfo shuffleTaskInfo = shuffleTaskInfos.get(appId);
- if (null == shuffleTaskInfo) {
+ // Once the shuffle server is restarted, the following shuffleTaskInfo may exist as
+ // a result of the refreshAppId() method being invoked during the process of sending/getting data
+ // from the client side. However, it is essential for the app's buffer pool to exist once it has been registered.
+ // This can serve as a crucial criterion for determining whether the app is not registered,
+ // especially following a server restart.
+ if (null == shuffleTaskInfos.get(appId) || !shuffleBufferManager.containsBufferByApp(appId)) {
  return RequireBufferStatusCode.NO_REGISTER.statusCode();
  }
+
  for (int partitionId : partitionIds) {
  long partitionUsedDataSize = getPartitionDataSize(appId, shuffleId, partitionId);
  if (shuffleBufferManager.limitHugePartition(

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
@@ -654,4 +654,8 @@ public boolean limitHugePartition(
  }
  return false;
  }
+
+ public boolean containsBufferByApp(String appId) {
+ return bufferPool.containsKey(appId);
+ }
 }