J-HowHuang commented on code in PR #15175:
URL: https://github.com/apache/pinot/pull/15175#discussion_r1994421697
##########
pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerClusterStatelessTest.java:
##########
@@ -513,6 +533,106 @@ public void testRebalance()
assertNull(rebalanceResult.getPreChecksResult());
_helixResourceManager.deleteOfflineTable(RAW_TABLE_NAME);
+
+ for (int i = 0; i < numServers; i++) {
+ stopAndDropFakeInstance(SERVER_INSTANCE_ID_PREFIX + i);
+ }
+ for (int i = 0; i < numServersToAdd; i++) {
+ stopAndDropFakeInstance(SERVER_INSTANCE_ID_PREFIX + (numServers + i));
+ }
+ executorService.shutdown();
+ }
+
+ @Test
+ public void testRebalancePreCheckerDiskUtil()
+ throws Exception {
+ int numServers = 3;
+ // Mock disk usage
+ Map<String, DiskUsageInfo> diskUsageInfoMap = new HashMap<>();
+
+ for (int i = 0; i < numServers; i++) {
+ String instanceId = "preCheckerDiskUtil_" + SERVER_INSTANCE_ID_PREFIX +
i;
+ addFakeServerInstanceToAutoJoinHelixCluster(instanceId, true);
+ DiskUsageInfo diskUsageInfo1 =
+ new DiskUsageInfo(instanceId, "", 1000L, 200L,
System.currentTimeMillis());
+ diskUsageInfoMap.put(instanceId, diskUsageInfo1);
+ }
+
+ ExecutorService executorService = Executors.newFixedThreadPool(10);
+ DefaultRebalancePreChecker preChecker = new DefaultRebalancePreChecker();
+ preChecker.init(_helixResourceManager, executorService, 0.5);
+ TableRebalancer tableRebalancer = new TableRebalancer(_helixManager, null,
null, preChecker,
+ _helixResourceManager.getTableSizeReader());
+ TableConfig tableConfig =
+ new
TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).setNumReplicas(NUM_REPLICAS).build();
+
+ // Create the table
+ addDummySchema(RAW_TABLE_NAME);
+ _helixResourceManager.addTable(tableConfig);
+
+ // Add the segments
+ int numSegments = 10;
+ for (int i = 0; i < numSegments; i++) {
+ _helixResourceManager.addNewSegment(OFFLINE_TABLE_NAME,
+ SegmentMetadataMockUtils.mockSegmentMetadata(RAW_TABLE_NAME,
SEGMENT_NAME_PREFIX + i), null);
+ }
+ Map<String, Map<String, String>> oldSegmentAssignment =
+
_helixResourceManager.getTableIdealState(OFFLINE_TABLE_NAME).getRecord().getMapFields();
+
+ // Add 3 more servers
+ int numServersToAdd = 3;
+ for (int i = 0; i < numServersToAdd; i++) {
+ String instanceId = "preCheckerDiskUtil_" + SERVER_INSTANCE_ID_PREFIX +
(numServers + i);
+ addFakeServerInstanceToAutoJoinHelixCluster(instanceId, true);
+ DiskUsageInfo diskUsageInfo =
+ new DiskUsageInfo(instanceId, "", 1000L, 200L,
System.currentTimeMillis());
+ diskUsageInfoMap.put(instanceId, diskUsageInfo);
+ }
+
+ ResourceUtilizationInfo.setDiskUsageInfo(diskUsageInfoMap);
+
+ // Rebalance in dry-run mode
+ RebalanceConfig rebalanceConfig = new RebalanceConfig();
+ rebalanceConfig.setDryRun(true);
+ rebalanceConfig.setPreChecks(true);
+
+ RebalanceResult rebalanceResult = tableRebalancer.rebalance(tableConfig,
rebalanceConfig, null);
+ assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE);
+ Map<String, String> preCheckResult = rebalanceResult.getPreChecksResult();
+ assertNotNull(preCheckResult);
+
assertTrue(preCheckResult.containsKey(DefaultRebalancePreChecker.DISK_UTILIZATION));
+ // Sending request to servers should fail for all, so needsPreprocess
should be set to "error" to indicate that a
Review Comment:
Thanks for the catch
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org
For additional commands, e-mail: commits-h...@pinot.apache.org
