desaikomal commented on code in PR #2543: URL: https://github.com/apache/helix/pull/2543#discussion_r1240865606
########## helix-core/src/test/java/org/apache/helix/integration/rebalancer/WagedRebalancer/TestWagedClusterExpansion.java: ########## @@ -0,0 +1,362 @@ +package org.apache.helix.integration.rebalancer.WagedRebalancer; + +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional warnrmation + * 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. + */ + +import java.util.ArrayList; +import java.util.Collections; +import java.util.Date; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import com.google.common.collect.ImmutableMap; +import org.apache.helix.HelixDataAccessor; +import org.apache.helix.NotificationContext; +import org.apache.helix.common.ZkTestBase; +import org.apache.helix.controller.rebalancer.waged.AssignmentMetadataStore; +import org.apache.helix.integration.manager.ClusterControllerManager; +import org.apache.helix.integration.manager.MockParticipantManager; +import org.apache.helix.manager.zk.ZKHelixDataAccessor; +import org.apache.helix.manager.zk.ZkBucketDataAccessor; +import org.apache.helix.model.BuiltInStateModelDefinitions; +import org.apache.helix.model.ClusterConfig; +import org.apache.helix.model.CurrentState; +import org.apache.helix.model.IdealState; +import org.apache.helix.model.LiveInstance; +import org.apache.helix.model.Message; +import org.apache.helix.model.ResourceAssignment; +import org.apache.helix.model.ResourceConfig; +import org.apache.helix.participant.StateMachineEngine; +import org.apache.helix.participant.statemachine.StateModel; +import org.apache.helix.participant.statemachine.StateModelFactory; +import org.apache.helix.participant.statemachine.StateModelInfo; +import org.apache.helix.participant.statemachine.Transition; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +/** + * This test case is specially targeting for n - n+1 issue. + * Waged is constraint-based algorithm. The end-state of calculated ideal-state will + * always satisfy hard-constraint. + * The issue is in order to reach the ideal-state, we need to do intermediate + * state-transitions. + * During this phase, hard-constraints are not satisfied. + * This test case is trying to mimic this by having a very tightly constraint + * cluster and increasing the weight for some resource and checking if + * intermediate state satisfies the need or not. + */ +public class TestWagedClusterExpansion extends ZkTestBase { + protected final int NUM_NODE = 6; + protected static final int START_PORT = 13000; + protected static final int PARTITIONS = 10; + + protected final String CLASS_NAME = getShortClassName(); + protected final String CLUSTER_NAME = CLUSTER_PREFIX + "_" + CLASS_NAME; + protected ClusterControllerManager _controller; + protected AssignmentMetadataStore _assignmentMetadataStore; + + List<MockParticipantManager> _participants = new ArrayList<>(); + + List<String> _nodes = new ArrayList<>(); + private final Set<String> _allDBs = new HashSet<>(); + private final int _replica = 3; + + private final String _testCapacityKey = "TestCapacityKey"; + private final Map<String, IdealState> _prevIdealState = new HashMap<>(); + + private static final Logger LOG = LoggerFactory.getLogger("TestWagedClusterExpansion"); + + + // mock delay master-slave state model + @StateModelInfo(initialState = "OFFLINE", states = { + "MASTER", "SLAVE", "ERROR" + }) + public static class WagedMasterSlaveModel extends StateModel { + private static final Logger LOG = LoggerFactory.getLogger(WagedMasterSlaveModel.class); + private final long _delay; + + public WagedMasterSlaveModel(long delay) { + _delay = delay; + } + + @Transition(to = "SLAVE", from = "OFFLINE") + public void onBecomeSlaveFromOffline(Message message, NotificationContext context) { + LOG.info("Become SLAVE from OFFLINE"); + } + + @Transition(to = "MASTER", from = "SLAVE") + public void onBecomeMasterFromSlave(Message message, NotificationContext context) + throws InterruptedException { + LOG.info("Become MASTER from SLAVE"); + } + + @Transition(to = "SLAVE", from = "MASTER") + public void onBecomeSlaveFromMaster(Message message, NotificationContext context) { + LOG.info("Become Slave from Master"); + } + + @Transition(to = "OFFLINE", from = "SLAVE") + public void onBecomeOfflineFromSlave(Message message, NotificationContext context) { + if (_delay > 0) { + try { + Thread.currentThread().sleep(_delay); + } catch (InterruptedException e) { + // ignore + } + } + LOG.info("Become OFFLINE from SLAVE"); + } + + @Transition(to = "DROPPED", from = "OFFLINE") + public void onBecomeDroppedFromOffline(Message message, NotificationContext context) { + if (_delay > 0) { + try { + Thread.currentThread().sleep(_delay); + } catch (InterruptedException e) { + // ignore + } + } + LOG.info("Become DROPPED FROM OFFLINE"); + } + } + + public class WagedDelayMSStateModelFactory extends StateModelFactory<WagedMasterSlaveModel> { + private long _delay; + + @Override + public WagedMasterSlaveModel createNewStateModel(String resourceName, + String partitionKey) { + WagedMasterSlaveModel model = new WagedMasterSlaveModel(_delay); + return model; + } + + public WagedDelayMSStateModelFactory setDelay(long delay) { + _delay = delay; + return this; + } + } + + @BeforeClass + public void beforeClass() throws Exception { + LOG.info("START " + CLASS_NAME + " at " + new Date(System.currentTimeMillis())); + + _gSetupTool.addCluster(CLUSTER_NAME, true); + // create 6 node cluster + for (int i = 0; i < NUM_NODE; i++) { + String storageNodeName = PARTICIPANT_PREFIX + "_" + (START_PORT + i); + _gSetupTool.addInstanceToCluster(CLUSTER_NAME, storageNodeName); + _nodes.add(storageNodeName); + } + // ST downward message will get delayed by 5sec. + startParticipants(5000); + + // start controller + String controllerName = CONTROLLER_PREFIX + "_0"; + _controller = new ClusterControllerManager(ZK_ADDR, CLUSTER_NAME, controllerName); + _controller.syncStart(); + + enablePersistBestPossibleAssignment(_gZkClient, CLUSTER_NAME, true); + + _assignmentMetadataStore = + new AssignmentMetadataStore(new ZkBucketDataAccessor(ZK_ADDR), CLUSTER_NAME) { + public Map<String, ResourceAssignment> getBaseline() { + // Ensure this metadata store always read from the ZK without using cache. + super.reset(); + return super.getBaseline(); + } + + public synchronized Map<String, ResourceAssignment> getBestPossibleAssignment() { + // Ensure this metadata store always read from the ZK without using cache. + super.reset(); + return super.getBestPossibleAssignment(); + } + }; + + // Set test instance capacity and partition weights + HelixDataAccessor dataAccessor = new ZKHelixDataAccessor(CLUSTER_NAME, _baseAccessor); + ClusterConfig clusterConfig = + dataAccessor.getProperty(dataAccessor.keyBuilder().clusterConfig()); + + clusterConfig.setInstanceCapacityKeys(Collections.singletonList(_testCapacityKey)); + clusterConfig.setDefaultInstanceCapacityMap(Collections.singletonMap(_testCapacityKey, 100)); + clusterConfig.setDefaultPartitionWeightMap(Collections.singletonMap(_testCapacityKey, 6)); + dataAccessor.setProperty(dataAccessor.keyBuilder().clusterConfig(), clusterConfig); + } + + private void startParticipants(int delay) { + // start dummy participants + for (String node : _nodes) { + MockParticipantManager participant = new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, node); + StateMachineEngine stateMach = participant.getStateMachineEngine(); + TestWagedClusterExpansion.WagedDelayMSStateModelFactory delayFactory = + new TestWagedClusterExpansion.WagedDelayMSStateModelFactory().setDelay(delay); + stateMach.registerStateModelFactory("MasterSlave", delayFactory); + participant.syncStart(); + _participants.add(participant); + } + } + + @Test + public void testInitialPlacement() { + // Create 3 resources with 10 partitions each. + for (int i = 0; i < 3; i++) { + String db = "Test-WagedDB-" + i; + createResourceWithWagedRebalance(CLUSTER_NAME, db, BuiltInStateModelDefinitions.MasterSlave.name(), + PARTITIONS, _replica, _replica); + _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, db, _replica); + _allDBs.add(db); + } + } + + @Test(dependsOnMethods = { "testInitialPlacement"}) + public void testIncreaseResourcePartitionWeight() throws Exception { + // For this test, let us record our initial prevIdealState. + for (String db : _allDBs) { + IdealState is = + _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, db); + _prevIdealState.put(db, is); + } + + // Let us add one more instance + String storageNodeName = PARTICIPANT_PREFIX + "_" + (START_PORT + NUM_NODE); + _gSetupTool.addInstanceToCluster(CLUSTER_NAME, storageNodeName); + _nodes.add(storageNodeName); + MockParticipantManager participant = new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, storageNodeName); + participant.syncStart(); + _participants.add(participant); + + // This is to make sure we have run the pipeline. + Thread.currentThread().sleep(2000); Review Comment: This is how most of the integration tests are written. I can invent better way to do it. -- 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: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
