[#2219] feat: Introduce ShuffleBlockIdManagerFactory and PartitionedShuffleBlockIdManager (#2227)

### What changes were proposed in this pull request?

- Introduce a ShuffleBlockIdManagerFactory to create configured implementation.
- Introduce PartitionedShuffleBlockIdManager to individual bitmap for each partition.

### Why are the changes needed?

Fix: #2219

<img width="1596" alt="image" src="https://github.com/user-attachments/assets/c7752f94-941b-45d2-9a45-1da197e38984">

This is the maximum partition num with node metrics in our cluster, it is ~410K, it cost heap size 410K * 50KiB = 2GiB if we say a bitmap related a partition avg 10000 blocks(cost 50KiB heap).

I think this is worth to choose this new policy in our production cluster.

### Does this PR introduce _any_ user-facing change?

- Config option: rss.server.blockIdStrategyClass

### How was this patch tested?

Test Locally, check result by `arthas`

- Common case, test with default config
```
[arthas@264]$ vmtool  --action getInstances --className org.apache.uniffle.server.ShuffleTaskInfo --express 'instances[0].shuffleBlockIdManager'
@DefaultShuffleBlockIdManager[
    LOG=@Log4jLogger[org.apache.logging.slf4j.Log4jLogger@3f211bae],
    partitionsToBlockIds=@ConcurrentHashMapForJDK8[isEmpty=false;size=1],
]

[arthas@264]$ vmtool  --action getInstances --className org.apache.uniffle.server.ShuffleTaskInfo --express 'instances[0].shuffleBlockIdManager.partitionsToBlockIds.get("application_1729845342052_431031_1730297736664") '
@ConcurrentHashMapForJDK8[
    @integer[0]:@Roaring64NavigableMap[][isEmpty=false;size=10],
    @integer[1]:@Roaring64NavigableMap[][isEmpty=false;size=10],
    @integer[2]:@Roaring64NavigableMap[][isEmpty=false;size=10],
]

```
- Config to PartitionedShuffleBlockIdManager
```
[arthas@264]$ vmtool  --action getInstances --className org.apache.uniffle.server.ShuffleTaskInfo --express 'instances[0].shuffleBlockIdManager'
@PartitionedShuffleBlockIdManager[
    LOG=@Log4jLogger[org.apache.logging.slf4j.Log4jLogger@6aa6fea0],
    partitionsToBlockIds=@ConcurrentHashMap[isEmpty=false;size=1],
]

[arthas@264]$ vmtool  --action getInstances --className org.apache.uniffle.server.ShuffleTaskInfo --express 'instances[0].shuffleBlockIdManager.partitionsToBlockIds'
@ConcurrentHashMap[
    @string[application_1729845342052_432324_1730299397916]:@ConcurrentHashMapForJDK8[isEmpty=false;size=2],
]
[arthas@264]$ vmtool  --action getInstances --className org.apache.uniffle.server.ShuffleTaskInfo --express 'instances[0].shuffleBlockIdManager.partitionsToBlockIds.get("application_1729845342052_432324_1730299397916")'
@ConcurrentHashMapForJDK8[
    @integer[0]:@ConcurrentHashMap[isEmpty=false;size=2000],
    @integer[1]:@ConcurrentHashMap[isEmpty=false;size=2000],
]
```
- Use client app level config
submit two app with `spark.rss.client.blockIdStrategyClass=org.apache.uniffle.server.block.DefaultShuffleBlockIdManager` and `spark.rss.client.blockIdStrategyClass=org.apache.uniffle.server.block.PartitionedShuffleBlockIdManager`

```
[arthas@264]$ vmtool  --action getInstances --className org.apache.uniffle.server.ShuffleTaskInfo --express 'instances[0].shuffleBlockIdManager'
@DefaultShuffleBlockIdManager[
    LOG=@Log4jLogger[org.apache.logging.slf4j.Log4jLogger@ed3b52a],
    partitionsToBlockIds=@ConcurrentHashMapForJDK8[isEmpty=false;size=1],
]
[arthas@264]$ vmtool  --action getInstances --className org.apache.uniffle.server.ShuffleTaskInfo --express 'instances[1].shuffleBlockIdManager'
@PartitionedShuffleBlockIdManager[
    LOG=@Log4jLogger[org.apache.logging.slf4j.Log4jLogger@5fc14628],
    partitionsToBlockIds=@ConcurrentHashMap[isEmpty=false;size=1],
]
```

Log example
```
[2024-10-30 23:01:37.585] [Grpc-709] [INFO] ShuffleTaskInfo - application_1729845342052_432387_1730300440568 use app configured ShuffleBlockIdManager to org.apache.uniffle.server.block.DefaultShuffleBlockIdManager@3c6b75a8
```

Author: maobaolong

common/src/main/java/org/apache/uniffle/common/config/RssClientConf.java

@@ -295,4 +295,12 @@ public class RssClientConf {
           .asList()
           .noDefaultValue()
           .withDescription("the extra java properties could be configured by this option");
+
+  public static final ConfigOption<String> RSS_CLIENT_BLOCK_ID_MANAGER_CLASS =
+      ConfigOptions.key("rss.client.blockIdManagerClass")
+          .stringType()
+          .noDefaultValue()
+          .withDescription(
+              "The block id manager class of server for this application, "
+                  + "the implementation of this interface to manage the shuffle block ids");
 }

docs/client_guide/client_guide.md

@@ -59,6 +59,7 @@ The important configuration of client is listed as following. These configuratio
 | <client_type>.rss.client.rpc.netty.pageSize                     | 4096                                   | The value of pageSize for PooledByteBufAllocator when using gRPC internal Netty on the client-side. This configuration will only take effect when rss.rpc.server.type is set to GRPC_NETTY.                                                                                                                                                                                                                                                                                                  |
 | <client_type>.rss.client.rpc.netty.maxOrder                     | 3                                      | The value of maxOrder for PooledByteBufAllocator when using gRPC internal Netty on the client-side. This configuration will only take effect when rss.rpc.server.type is set to GRPC_NETTY.                                                                                                                                                                                                                                                                                                  |
 | <client_type>.rss.client.rpc.netty.smallCacheSize               | 1024                                   | The value of smallCacheSize for PooledByteBufAllocator when using gRPC internal Netty on the client-side. This configuration will only take effect when rss.rpc.server.type is set to GRPC_NETTY.                                                                                                                                                                                                                                                                                            |
+| <client_type>.rss.client.blockIdManagerClass                    | -                                      | The block id manager class of server for this application, the implementation of this interface to manage the shuffle block ids                                                                                                                                                                                                                                                                                                                                                              |
 
 Notice:
 

docs/server_guide.md

@@ -121,6 +121,7 @@ This document will introduce how to deploy Uniffle shuffle servers.
 | rss.storage.localFileWriterClass                         | org.apache.uniffle.storage.handler.impl.LocalFileWriter                | The writer class to write shuffle data for LOCALFILE.                                                                                                                                                                                                                                                                                                                                                                                                                                                                             |
 | rss.storage.hdfs.write.dataBufferSize                    | 8K                                                                     | The size of the buffer used to cache data written for HDFS.                                                                                                                                                                                                                                                                                                                                                                                                                                                                       |
 | rss.storage.hdfs.write.indexBufferSize                   | 8K                                                                     | The size of the buffer used to cache index written for HDFS.                                                                                                                                                                                                                                                                                                                                                                                                                                                                      |
+| rss.server.blockIdManagerClass                           | org.apache.uniffle.server.block.DefaultShuffleBlockIdManager           | The block id manager class. It is used to manage block id.                                                                                                                                                                                                                                                                                                                                                                                                                                                                        |
 
 ### Advanced Configurations
 | Property Name                                    | Default | Description                                                                                                                                                                                 |

integration-test/common/src/test/java/org/apache/uniffle/test/ShuffleServerGrpcTest.java

@@ -352,14 +352,13 @@ public void shuffleResultTest() throws Exception {
     request = new RssReportShuffleResultRequest("shuffleResultTest", 2, 1L, partitionToBlockIds, 3);
     grpcShuffleServerClient.reportShuffleResult(request);
     // validate bitmap in shuffleTaskManager
-    Roaring64NavigableMap[] bitmaps =
+    long bitmapNum =
         grpcShuffleServers
             .get(0)
             .getShuffleTaskManager()
-            .getPartitionsToBlockIds()
-            .get("shuffleResultTest")
-            .get(2);
-    assertEquals(3, bitmaps.length);
+            .getShuffleBlockIdManager()
+            .getBitmapNum("shuffleResultTest", 2);
+    assertEquals(3, bitmapNum);
 
     req = new RssGetShuffleResultRequest("shuffleResultTest", 2, 1, layout);
     result = grpcShuffleServerClient.getShuffleResult(req);

integration-test/common/src/test/java/org/apache/uniffle/test/ShuffleServerWithLocalTest.java

@@ -179,10 +179,18 @@ private void localWriteReadTest(boolean isNettyMode) throws Exception {
 
     List<ShuffleServer> shuffleServers = isNettyMode ? nettyShuffleServers : grpcShuffleServers;
     assertNotNull(
-        shuffleServers.get(0).getShuffleTaskManager().getPartitionsToBlockIds().get(testAppId));
+        shuffleServers
+            .get(0)
+            .getShuffleTaskManager()
+            .getShuffleBlockIdManager()
+            .contains(testAppId));
     Thread.sleep(8000);
     assertNull(
-        shuffleServers.get(0).getShuffleTaskManager().getPartitionsToBlockIds().get(testAppId));
+        shuffleServers
+            .get(0)
+            .getShuffleTaskManager()
+            .getShuffleBlockIdManager()
+            .contains(testAppId));
   }
 
   protected void validateResult(

server/src/main/java/org/apache/uniffle/server/ShuffleServerConf.java

@@ -25,6 +25,7 @@
 import org.apache.uniffle.common.config.ConfigOptions;
 import org.apache.uniffle.common.config.ConfigUtils;
 import org.apache.uniffle.common.config.RssBaseConf;
+import org.apache.uniffle.server.block.DefaultShuffleBlockIdManager;
 import org.apache.uniffle.server.buffer.ShuffleBufferType;
 
 public class ShuffleServerConf extends RssBaseConf {
@@ -743,6 +744,13 @@ public class ShuffleServerConf extends RssBaseConf {
           .booleanType()
           .defaultValue(false)
           .withDescription("Whether to enable app detail log");
+  public static final ConfigOption<String> SERVER_BLOCK_ID_MANAGER_CLASS =
+      ConfigOptions.key("rss.server.blockIdManagerClass")
+          .stringType()
+          .defaultValue(DefaultShuffleBlockIdManager.class.getName())
+          .withDescription(
+              "The block id manager class, the implementation of this interface "
+                  + "to manage the shuffle block ids");
 
   public ShuffleServerConf() {}
 

server/src/main/java/org/apache/uniffle/server/ShuffleTaskInfo.java

@@ -26,14 +26,19 @@
 import java.util.stream.Collectors;
 
 import com.google.common.collect.Sets;
+import org.apache.commons.lang3.StringUtils;
 import org.roaringbitmap.longlong.Roaring64NavigableMap;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import org.apache.uniffle.common.PartitionInfo;
 import org.apache.uniffle.common.ShuffleDataDistributionType;
+import org.apache.uniffle.common.config.RssClientConf;
+import org.apache.uniffle.common.util.Constants;
 import org.apache.uniffle.common.util.JavaUtils;
 import org.apache.uniffle.common.util.UnitConverter;
+import org.apache.uniffle.server.block.ShuffleBlockIdManager;
+import org.apache.uniffle.server.block.ShuffleBlockIdManagerFactory;
 
 /**
  * ShuffleTaskInfo contains the information of submitting the shuffle, the information of the cache
@@ -78,6 +83,7 @@ public class ShuffleTaskInfo {
 
   private final Map<Integer, Integer> latestStageAttemptNumbers;
   private Map<String, String> properties;
+  private ShuffleBlockIdManager shuffleBlockIdManager;
 
   public ShuffleTaskInfo(String appId) {
     this.appId = appId;
@@ -324,6 +330,32 @@ public void setProperties(Map<String, String> properties) {
             .filter(entry -> entry.getKey().contains(".rss."))
             .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
     this.properties = filteredProperties;
-    LOGGER.info("{} set properties to {}", appId, properties);
+    LOGGER.info("{} set properties to {}", appId, filteredProperties);
+    String keyName = RssClientConf.RSS_CLIENT_BLOCK_ID_MANAGER_CLASS.key();
+    String className = properties.get(keyName);
+    if (StringUtils.isEmpty(className)) {
+      keyName =
+          Constants.SPARK_RSS_CONFIG_PREFIX + RssClientConf.RSS_CLIENT_BLOCK_ID_MANAGER_CLASS.key();
+      className =
+          properties.get(
+              Constants.SPARK_RSS_CONFIG_PREFIX
+                  + RssClientConf.RSS_CLIENT_BLOCK_ID_MANAGER_CLASS.key());
+    }
+    if (StringUtils.isNotEmpty(className)) {
+      shuffleBlockIdManager =
+          ShuffleBlockIdManagerFactory.createShuffleBlockIdManager(className, keyName);
+      LOGGER.info(
+          "{} use app configured ShuffleBlockIdManager to {}", appId, shuffleBlockIdManager);
+    }
+  }
+
+  public ShuffleBlockIdManager getShuffleBlockIdManager() {
+    return shuffleBlockIdManager;
+  }
+
+  public void setShuffleBlockIdManagerIfNeeded(ShuffleBlockIdManager shuffleBlockIdManager) {
+    if (this.shuffleBlockIdManager == null) {
+      this.shuffleBlockIdManager = shuffleBlockIdManager;
+    }
   }
 }