somandal commented on code in PR #15575:
URL: https://github.com/apache/pinot/pull/15575#discussion_r2054954239


##########
pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/TableRebalanceIntegrationTest.java:
##########
@@ -0,0 +1,1170 @@
+/**
+ * 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.pinot.integration.tests;
+
+import java.io.IOException;
+import java.net.URISyntaxException;
+import java.net.URL;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.pinot.common.exception.HttpErrorStatusException;
+import org.apache.pinot.common.utils.SimpleHttpResponse;
+import org.apache.pinot.common.utils.config.TagNameUtils;
+import org.apache.pinot.common.utils.http.HttpClient;
+import org.apache.pinot.common.utils.regex.JavaUtilPattern;
+import org.apache.pinot.common.utils.regex.Matcher;
+import org.apache.pinot.common.utils.regex.Pattern;
+import 
org.apache.pinot.controller.api.resources.ServerRebalanceJobStatusResponse;
+import 
org.apache.pinot.controller.api.resources.ServerReloadControllerJobStatusResponse;
+import 
org.apache.pinot.controller.helix.core.rebalance.DefaultRebalancePreChecker;
+import org.apache.pinot.controller.helix.core.rebalance.RebalanceConfig;
+import 
org.apache.pinot.controller.helix.core.rebalance.RebalancePreCheckerResult;
+import org.apache.pinot.controller.helix.core.rebalance.RebalanceResult;
+import org.apache.pinot.controller.helix.core.rebalance.RebalanceSummaryResult;
+import org.apache.pinot.server.starter.helix.BaseServerStarter;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.config.table.TagOverrideConfig;
+import org.apache.pinot.spi.config.table.TenantConfig;
+import org.apache.pinot.spi.config.table.assignment.InstanceAssignmentConfig;
+import org.apache.pinot.spi.config.table.assignment.InstanceConstraintConfig;
+import 
org.apache.pinot.spi.config.table.assignment.InstanceReplicaGroupPartitionConfig;
+import org.apache.pinot.spi.config.table.assignment.InstanceTagPoolConfig;
+import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.MetricFieldSpec;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.utils.JsonUtils;
+import org.apache.pinot.spi.utils.StringUtil;
+import org.apache.pinot.util.TestUtils;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertNull;
+import static org.testng.Assert.assertTrue;
+
+
+public class TableRebalanceIntegrationTest extends 
HybridClusterIntegrationTest {
+  private String getRebalanceUrl(RebalanceConfig rebalanceConfig, TableType 
tableType) {
+    return StringUtil.join("/", getControllerRequestURLBuilder().getBaseUrl(), 
"tables", getTableName(), "rebalance")
+        + "?type=" + tableType.toString() + "&" + 
rebalanceConfig.toQueryString();
+  }
+
+  @Test
+  public void testOfflineRebalancePreChecks()
+      throws Exception {
+    // setup the rebalance config
+    RebalanceConfig rebalanceConfig = new RebalanceConfig();
+    rebalanceConfig.setDryRun(true);
+
+    TableConfig tableConfig = getOfflineTableConfig();
+    TableConfig originalTableConfig = new TableConfig(tableConfig);
+
+    // Ensure pre-check status is null if not enabled
+    String response = sendPostRequest(getRebalanceUrl(rebalanceConfig, 
TableType.OFFLINE));
+    RebalanceResult rebalanceResult = JsonUtils.stringToObject(response, 
RebalanceResult.class);
+    assertNull(rebalanceResult.getPreChecksResult());
+
+    // Enable pre-checks, nothing is set
+    rebalanceConfig.setPreChecks(true);
+    response = sendPostRequest(getRebalanceUrl(rebalanceConfig, 
TableType.OFFLINE));
+    rebalanceResult = JsonUtils.stringToObject(response, 
RebalanceResult.class);
+    checkRebalancePreCheckStatus(rebalanceResult, RebalanceResult.Status.NO_OP,
+        "Instance assignment not allowed, no need for minimizeDataMovement",
+        RebalancePreCheckerResult.PreCheckStatus.PASS, "No need to reload",
+        RebalancePreCheckerResult.PreCheckStatus.PASS, "All rebalance 
parameters look good",
+        RebalancePreCheckerResult.PreCheckStatus.PASS,
+        "OFFLINE segments - Replica Groups are not enabled, replication: " + 
tableConfig.getReplication(),
+        RebalancePreCheckerResult.PreCheckStatus.PASS);
+
+    // Enable minimizeDataMovement
+    Map<String, InstanceAssignmentConfig> instanceAssignmentConfigMap =
+        Collections.singletonMap("OFFLINE", 
createInstanceAssignmentConfig(true, TableType.OFFLINE));
+    InstanceReplicaGroupPartitionConfig replicaGroupPartitionConfig =
+        
instanceAssignmentConfigMap.get("OFFLINE").getReplicaGroupPartitionConfig();
+    tableConfig.setInstanceAssignmentConfigMap(instanceAssignmentConfigMap);
+    updateTableConfig(tableConfig);
+    response = sendPostRequest(getRebalanceUrl(rebalanceConfig, 
TableType.OFFLINE));
+    rebalanceResult = JsonUtils.stringToObject(response, 
RebalanceResult.class);
+    checkRebalancePreCheckStatus(rebalanceResult, RebalanceResult.Status.NO_OP,
+        "minimizeDataMovement is enabled", 
RebalancePreCheckerResult.PreCheckStatus.PASS,
+        "No need to reload", RebalancePreCheckerResult.PreCheckStatus.PASS,
+        "All rebalance parameters look good",
+        RebalancePreCheckerResult.PreCheckStatus.PASS,
+        "reassignInstances is disabled, replica groups may not be 
updated.\nOFFLINE segments - numReplicaGroups: "
+            + replicaGroupPartitionConfig.getNumReplicaGroups() + ", 
numInstancesPerReplicaGroup: "
+            + (replicaGroupPartitionConfig.getNumInstancesPerReplicaGroup() == 0
+            ? "0 (using as many instances as possible)" : 
replicaGroupPartitionConfig.getNumInstancesPerReplicaGroup()),
+        RebalancePreCheckerResult.PreCheckStatus.WARN);
+
+    // Override minimizeDataMovement
+    
rebalanceConfig.setMinimizeDataMovement(RebalanceConfig.MinimizeDataMovementOptions.DISABLE);
+    response = sendPostRequest(getRebalanceUrl(rebalanceConfig, 
TableType.OFFLINE));
+    rebalanceResult = JsonUtils.stringToObject(response, 
RebalanceResult.class);
+    checkRebalancePreCheckStatus(rebalanceResult, RebalanceResult.Status.NO_OP,
+        "minimizeDataMovement is enabled in table config but it's overridden 
with disabled",
+        RebalancePreCheckerResult.PreCheckStatus.WARN, "No need to reload",
+        RebalancePreCheckerResult.PreCheckStatus.PASS, "All rebalance 
parameters look good",
+        RebalancePreCheckerResult.PreCheckStatus.PASS,
+        "reassignInstances is disabled, replica groups may not be 
updated.\nOFFLINE segments - numReplicaGroups: "
+            + replicaGroupPartitionConfig.getNumReplicaGroups() + ", 
numInstancesPerReplicaGroup: "
+            + (replicaGroupPartitionConfig.getNumInstancesPerReplicaGroup() == 0
+            ? "0 (using as many instances as possible)" : 
replicaGroupPartitionConfig.getNumInstancesPerReplicaGroup()),
+        RebalancePreCheckerResult.PreCheckStatus.WARN);
+
+    // Use default minimizeDataMovement and disable it in table config
+    instanceAssignmentConfigMap =
+        Collections.singletonMap("OFFLINE", 
createInstanceAssignmentConfig(false, TableType.OFFLINE));
+    replicaGroupPartitionConfig = 
instanceAssignmentConfigMap.get("OFFLINE").getReplicaGroupPartitionConfig();
+    tableConfig.setInstanceAssignmentConfigMap(instanceAssignmentConfigMap);
+    updateTableConfig(tableConfig);
+    
rebalanceConfig.setMinimizeDataMovement(RebalanceConfig.MinimizeDataMovementOptions.DEFAULT);
+    response = sendPostRequest(getRebalanceUrl(rebalanceConfig, 
TableType.OFFLINE));
+    rebalanceResult = JsonUtils.stringToObject(response, 
RebalanceResult.class);
+    checkRebalancePreCheckStatus(rebalanceResult, RebalanceResult.Status.NO_OP,
+        "minimizeDataMovement is not enabled but instance assignment is 
allowed",
+        RebalancePreCheckerResult.PreCheckStatus.WARN, "No need to reload",
+        RebalancePreCheckerResult.PreCheckStatus.PASS, "All rebalance 
parameters look good",
+        RebalancePreCheckerResult.PreCheckStatus.PASS,
+        "reassignInstances is disabled, replica groups may not be 
updated.\nOFFLINE segments - numReplicaGroups: "
+            + replicaGroupPartitionConfig.getNumReplicaGroups() + ", 
numInstancesPerReplicaGroup: "
+            + (replicaGroupPartitionConfig.getNumInstancesPerReplicaGroup() == 0
+            ? "0 (using as many instances as possible)" : 
replicaGroupPartitionConfig.getNumInstancesPerReplicaGroup()),
+        RebalancePreCheckerResult.PreCheckStatus.WARN);
+
+    // Undo minimizeDataMovement, update the table config to add a column to 
bloom filter
+    
rebalanceConfig.setMinimizeDataMovement(RebalanceConfig.MinimizeDataMovementOptions.ENABLE);
+    tableConfig.getIndexingConfig().getBloomFilterColumns().add("Quarter");
+    tableConfig.setInstanceAssignmentConfigMap(null);
+    updateTableConfig(tableConfig);
+    response = sendPostRequest(getRebalanceUrl(rebalanceConfig, 
TableType.OFFLINE));
+    rebalanceResult = JsonUtils.stringToObject(response, 
RebalanceResult.class);
+    checkRebalancePreCheckStatus(rebalanceResult, RebalanceResult.Status.NO_OP,
+        "Instance assignment not allowed, no need for minimizeDataMovement",
+        RebalancePreCheckerResult.PreCheckStatus.PASS, "Reload needed prior to 
running rebalance",
+        RebalancePreCheckerResult.PreCheckStatus.WARN, "All rebalance 
parameters look good",
+        RebalancePreCheckerResult.PreCheckStatus.PASS,
+        "OFFLINE segments - Replica Groups are not enabled, replication: " + 
tableConfig.getReplication(),
+        RebalancePreCheckerResult.PreCheckStatus.PASS);
+
+    // Undo tableConfig change
+    tableConfig.getIndexingConfig().getBloomFilterColumns().remove("Quarter");
+    updateTableConfig(tableConfig);
+    response = sendPostRequest(getRebalanceUrl(rebalanceConfig, 
TableType.OFFLINE));
+    rebalanceResult = JsonUtils.stringToObject(response, 
RebalanceResult.class);
+    checkRebalancePreCheckStatus(rebalanceResult, RebalanceResult.Status.NO_OP,
+        "Instance assignment not allowed, no need for minimizeDataMovement",
+        RebalancePreCheckerResult.PreCheckStatus.PASS, "No need to reload",
+        RebalancePreCheckerResult.PreCheckStatus.PASS, "All rebalance 
parameters look good",
+        RebalancePreCheckerResult.PreCheckStatus.PASS,
+        "OFFLINE segments - Replica Groups are not enabled, replication: " + 
tableConfig.getReplication(),
+        RebalancePreCheckerResult.PreCheckStatus.PASS);
+
+    // Add a new server (to force change in instance assignment) and enable 
reassignInstances
+    // Validate that the status for reload is still PASS (i.e. even though an 
extra server is tagged which has no
+    // segments assigned for this table, we don't try to get needReload status 
from that extra server, otherwise
+    // ERROR status would be returned)
+    BaseServerStarter serverStarter0 = startOneServer(NUM_SERVERS);
+    createServerTenant(getServerTenant(), 1, 0);
+    rebalanceConfig.setReassignInstances(true);
+    tableConfig.setInstanceAssignmentConfigMap(null);
+    updateTableConfig(tableConfig);
+    response = sendPostRequest(getRebalanceUrl(rebalanceConfig, 
TableType.OFFLINE));
+    rebalanceResult = JsonUtils.stringToObject(response, 
RebalanceResult.class);
+    checkRebalancePreCheckStatus(rebalanceResult, RebalanceResult.Status.DONE,
+        "Instance assignment not allowed, no need for minimizeDataMovement",
+        RebalancePreCheckerResult.PreCheckStatus.PASS, "No need to reload",
+        RebalancePreCheckerResult.PreCheckStatus.PASS, "All rebalance 
parameters look good",
+        RebalancePreCheckerResult.PreCheckStatus.PASS,
+        "OFFLINE segments - Replica Groups are not enabled, replication: " + 
tableConfig.getReplication(),
+        RebalancePreCheckerResult.PreCheckStatus.PASS);
+    rebalanceConfig.setReassignInstances(false);
+
+    // Stop the added server
+    serverStarter0.stop();
+    TestUtils.waitForCondition(
+        aVoid -> 
getHelixResourceManager().dropInstance(serverStarter0.getInstanceId()).isSuccessful(),
+        60_000L, "Failed to drop added server");
+
+    // Add a schema change. Notice that this may affect other following tests
+    Schema schema = createSchema();
+    schema.addField(new MetricFieldSpec("NewAddedIntMetricA", 
FieldSpec.DataType.INT, 1));
+    updateSchema(schema);
+    response = sendPostRequest(getRebalanceUrl(rebalanceConfig, 
TableType.OFFLINE));
+    rebalanceResult = JsonUtils.stringToObject(response, 
RebalanceResult.class);
+    checkRebalancePreCheckStatus(rebalanceResult, RebalanceResult.Status.NO_OP,
+        "Instance assignment not allowed, no need for minimizeDataMovement",
+        RebalancePreCheckerResult.PreCheckStatus.PASS, "Reload needed prior to 
running rebalance",
+        RebalancePreCheckerResult.PreCheckStatus.WARN, "All rebalance 
parameters look good",
+        RebalancePreCheckerResult.PreCheckStatus.PASS,
+        "OFFLINE segments - Replica Groups are not enabled, replication: " + 
tableConfig.getReplication(),
+        RebalancePreCheckerResult.PreCheckStatus.PASS);
+
+    // Keep schema change and update table config to add minimizeDataMovement
+    instanceAssignmentConfigMap =
+        Collections.singletonMap("OFFLINE", 
createInstanceAssignmentConfig(true, TableType.OFFLINE));
+    replicaGroupPartitionConfig = 
instanceAssignmentConfigMap.get("OFFLINE").getReplicaGroupPartitionConfig();
+    tableConfig.setInstanceAssignmentConfigMap(instanceAssignmentConfigMap);
+    updateTableConfig(tableConfig);
+    response = sendPostRequest(getRebalanceUrl(rebalanceConfig, 
TableType.OFFLINE));
+    rebalanceResult = JsonUtils.stringToObject(response, 
RebalanceResult.class);
+    checkRebalancePreCheckStatus(rebalanceResult, RebalanceResult.Status.NO_OP,
+        "minimizeDataMovement is enabled", 
RebalancePreCheckerResult.PreCheckStatus.PASS,
+        "Reload needed prior to running rebalance", 
RebalancePreCheckerResult.PreCheckStatus.WARN,
+        "All rebalance parameters look good",
+        RebalancePreCheckerResult.PreCheckStatus.PASS,
+        "reassignInstances is disabled, replica groups may not be 
updated.\nOFFLINE segments - numReplicaGroups: "
+            + replicaGroupPartitionConfig.getNumReplicaGroups() + ", 
numInstancesPerReplicaGroup: "
+            + (replicaGroupPartitionConfig.getNumInstancesPerReplicaGroup() == 0
+            ? "0 (using as many instances as possible)" : 
replicaGroupPartitionConfig.getNumInstancesPerReplicaGroup()),
+        RebalancePreCheckerResult.PreCheckStatus.WARN);
+
+    // Keep schema change and update table config to add instance config map 
with minimizeDataMovement = false
+    instanceAssignmentConfigMap =
+        Collections.singletonMap("OFFLINE", 
createInstanceAssignmentConfig(false, TableType.OFFLINE));
+    replicaGroupPartitionConfig = 
instanceAssignmentConfigMap.get("OFFLINE").getReplicaGroupPartitionConfig();
+    tableConfig.setInstanceAssignmentConfigMap(instanceAssignmentConfigMap);
+    updateTableConfig(tableConfig);
+    response = sendPostRequest(getRebalanceUrl(rebalanceConfig, 
TableType.OFFLINE));
+    rebalanceResult = JsonUtils.stringToObject(response, 
RebalanceResult.class);
+    checkRebalancePreCheckStatus(rebalanceResult, RebalanceResult.Status.NO_OP,
+        "minimizeDataMovement is enabled",
+        RebalancePreCheckerResult.PreCheckStatus.PASS, "Reload needed prior to 
running rebalance",
+        RebalancePreCheckerResult.PreCheckStatus.WARN, "All rebalance 
parameters look good",
+        RebalancePreCheckerResult.PreCheckStatus.PASS,
+        "reassignInstances is disabled, replica groups may not be 
updated.\nOFFLINE segments - numReplicaGroups: "
+            + replicaGroupPartitionConfig.getNumReplicaGroups() + ", 
numInstancesPerReplicaGroup: "
+            + (replicaGroupPartitionConfig.getNumInstancesPerReplicaGroup() == 0
+            ? "0 (using as many instances as possible)" : 
replicaGroupPartitionConfig.getNumInstancesPerReplicaGroup()),
+        RebalancePreCheckerResult.PreCheckStatus.WARN);
+
+    // Add a new server (to force change in instance assignment) and enable 
reassignInstances
+    // Trigger rebalance config warning
+    BaseServerStarter serverStarter1 = startOneServer(NUM_SERVERS + 1);
+    createServerTenant(getServerTenant(), 1, 0);
+    rebalanceConfig.setReassignInstances(true);
+    rebalanceConfig.setBestEfforts(true);
+    rebalanceConfig.setBootstrap(true);
+    rebalanceConfig.setMinAvailableReplicas(-1);
+    tableConfig.setInstanceAssignmentConfigMap(null);
+    updateTableConfig(tableConfig);
+    response = sendPostRequest(getRebalanceUrl(rebalanceConfig, 
TableType.OFFLINE));
+    rebalanceResult = JsonUtils.stringToObject(response, 
RebalanceResult.class);
+    checkRebalancePreCheckStatus(rebalanceResult, RebalanceResult.Status.DONE,
+        "Instance assignment not allowed, no need for minimizeDataMovement",
+        RebalancePreCheckerResult.PreCheckStatus.PASS, "Reload needed prior to 
running rebalance",
+        RebalancePreCheckerResult.PreCheckStatus.WARN,
+        "bestEfforts is enabled, only enable it if you know what you are 
doing\n"
+            + "bootstrap is enabled which can cause a large amount of data 
movement, double check if this is "
+            + "intended", RebalancePreCheckerResult.PreCheckStatus.WARN,
+        "OFFLINE segments - Replica Groups are not enabled, replication: " + 
tableConfig.getReplication(),
+        RebalancePreCheckerResult.PreCheckStatus.PASS);
+
+    // reload
+    response =
+        
sendPostRequest(getControllerRequestURLBuilder().forTableReload(getTableName(), 
TableType.OFFLINE, true));
+    waitForReloadToComplete(getReloadJobIdFromResponse(response), 30_000);
+    // reload realtime table as well for other realtime tests
+    response =
+        
sendPostRequest(getControllerRequestURLBuilder().forTableReload(getTableName(), 
TableType.REALTIME, false));
+    waitForReloadToComplete(getReloadJobIdFromResponse(response), 30_000);
+    response = sendPostRequest(getRebalanceUrl(rebalanceConfig, 
TableType.OFFLINE));
+    rebalanceResult = JsonUtils.stringToObject(response, 
RebalanceResult.class);
+    checkRebalancePreCheckStatus(rebalanceResult, RebalanceResult.Status.DONE,
+        "Instance assignment not allowed, no need for minimizeDataMovement",
+        RebalancePreCheckerResult.PreCheckStatus.PASS, "No need to reload",
+        RebalancePreCheckerResult.PreCheckStatus.PASS,
+        "bestEfforts is enabled, only enable it if you know what you are 
doing\n"
+            + "bootstrap is enabled which can cause a large amount of data 
movement, double check if this is "
+            + "intended", RebalancePreCheckerResult.PreCheckStatus.WARN,
+        "OFFLINE segments - Replica Groups are not enabled, replication: " + 
tableConfig.getReplication(),
+        RebalancePreCheckerResult.PreCheckStatus.PASS);
+
+    // Disable dry-run
+    rebalanceConfig.setBootstrap(false);
+    rebalanceConfig.setBestEfforts(false);
+    rebalanceConfig.setDryRun(false);
+    response = sendPostRequest(getRebalanceUrl(rebalanceConfig, 
TableType.OFFLINE));
+    rebalanceResult = JsonUtils.stringToObject(response, 
RebalanceResult.class);
+    assertNull(rebalanceResult.getPreChecksResult());
+    assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.FAILED);

Review Comment:
   got it, thanks for digging in. i should've added a better comment lol



-- 
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

Reply via email to