(incubator-uniffle) branch master updated: [MINOR] fix(server): Assert actual number of bitmaps matches bitNum (#1493)

[zuston]

This is an automated email from the ASF dual-hosted git repository.

zuston pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-uniffle.git


The following commit(s) were added to refs/heads/master by this push:
     new bd5007b08 [MINOR] fix(server): Assert actual number of bitmaps matches 
bitNum (#1493)
bd5007b08 is described below

commit bd5007b08f9107ee2e3e3e64c8e5a219c67bf021
Author: Enrico Minack <git...@enrico.minack.dev>
AuthorDate: Tue Jan 30 04:12:04 2024 +0100

    [MINOR] fix(server): Assert actual number of bitmaps matches bitNum (#1493)
    
    ### What changes were proposed in this pull request?
    Method `ShuffleTaskManager.addFinishedBlockIds` should check the actual 
number of bitmaps matches the requested `bitmapNum`.
    
    ### Why are the changes needed?
    Calling `ShuffleTaskManager.addFinishedBlockIds(String appId, Integer 
shuffleId, …, bitmapNum)` initializes an array of bitmaps for the first call 
per `(appId, shuffleId)`, according to the `bitmapNum`. Subsequent calls must 
use the same `bitmapNum`, otherwise you get into trouble:
    
    - with a larger `bimapNum` you would see an out-of-bound exception, e.g.: 
`java.lang.ArrayIndexOutOfBoundsException: Index 3 out of bounds for length 3`
    - a lower `bimapNum` would silently add the blocks to the wrong bitmap, 
potentially hiding them from shuffle read clients
    
    ### Does this PR introduce _any_ user-facing change?
    No.
    
    ### How was this patch tested?
    Test added to `ShuffleTaskManagerTest.testGetFinishedBlockIds`.
---
 .../java/org/apache/uniffle/server/ShuffleTaskManager.java     | 10 ++++++++++
 .../java/org/apache/uniffle/server/ShuffleTaskManagerTest.java |  9 +++++++++
 2 files changed, 19 insertions(+)

diff --git 
a/server/src/main/java/org/apache/uniffle/server/ShuffleTaskManager.java 
b/server/src/main/java/org/apache/uniffle/server/ShuffleTaskManager.java
index efcd56150..04d62dabe 100644
--- a/server/src/main/java/org/apache/uniffle/server/ShuffleTaskManager.java
+++ b/server/src/main/java/org/apache/uniffle/server/ShuffleTaskManager.java
@@ -56,6 +56,7 @@ import org.apache.uniffle.common.ShufflePartitionedBlock;
 import org.apache.uniffle.common.ShufflePartitionedData;
 import org.apache.uniffle.common.config.RssBaseConf;
 import org.apache.uniffle.common.exception.FileNotFoundException;
+import org.apache.uniffle.common.exception.InvalidRequestException;
 import org.apache.uniffle.common.exception.NoBufferException;
 import org.apache.uniffle.common.exception.NoBufferForHugePartitionException;
 import org.apache.uniffle.common.exception.NoRegisterException;
@@ -386,6 +387,15 @@ public class ShuffleTaskManager {
           return blockIds;
         });
     Roaring64NavigableMap[] blockIds = shuffleIdToPartitions.get(shuffleId);
+    if (blockIds.length != bitmapNum) {
+      throw new InvalidRequestException(
+          "Request expects "
+              + bitmapNum
+              + " bitmaps, but there are "
+              + blockIds.length
+              + " bitmaps!");
+    }
+
     for (Map.Entry<Integer, long[]> entry : partitionToBlockIds.entrySet()) {
       Integer partitionId = entry.getKey();
       Roaring64NavigableMap bitmap = blockIds[partitionId % bitmapNum];
diff --git 
a/server/src/test/java/org/apache/uniffle/server/ShuffleTaskManagerTest.java 
b/server/src/test/java/org/apache/uniffle/server/ShuffleTaskManagerTest.java
index b0999bce1..fb1de11c2 100644
--- a/server/src/test/java/org/apache/uniffle/server/ShuffleTaskManagerTest.java
+++ b/server/src/test/java/org/apache/uniffle/server/ShuffleTaskManagerTest.java
@@ -49,6 +49,7 @@ import org.apache.uniffle.common.RemoteStorageInfo;
 import org.apache.uniffle.common.ShuffleDataResult;
 import org.apache.uniffle.common.ShufflePartitionedBlock;
 import org.apache.uniffle.common.ShufflePartitionedData;
+import org.apache.uniffle.common.exception.InvalidRequestException;
 import org.apache.uniffle.common.exception.NoBufferForHugePartitionException;
 import org.apache.uniffle.common.exception.NoRegisterException;
 import org.apache.uniffle.common.exception.RssException;
@@ -851,6 +852,14 @@ public class ShuffleTaskManagerTest extends HadoopTestBase 
{
         shuffleTaskManager.getFinishedBlockIds(appId, shuffleId, 
requestPartitions);
     Roaring64NavigableMap resBlockIds = 
RssUtils.deserializeBitMap(serializeBitMap);
     assertEquals(expectedBlockIds, resBlockIds);
+
+    try {
+      // calling with same appId and shuffleId but different bitmapNum should 
fail
+      shuffleTaskManager.addFinishedBlockIds(appId, shuffleId, 
blockIdsToReport, bitNum - 1);
+      fail("Exception should be thrown");
+    } catch (InvalidRequestException e) {
+      assertEquals(e.getMessage(), "Request expects 2 bitmaps, but there are 3 
bitmaps!");
+    }
   }
 
   @Test