helix-core/src/test/java/org/apache/helix/integration/TestWeightBasedRebalanceUtil.java - helix - Git at Google

 package org.apache.helix.integration;

 /*
  * 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.
  */

 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Date;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;

 import org.apache.helix.HelixException;
 import org.apache.helix.zookeeper.datamodel.ZNRecord;
 import org.apache.helix.api.config.RebalanceConfig;
 import org.apache.helix.api.rebalancer.constraint.AbstractRebalanceHardConstraint;
 import org.apache.helix.api.rebalancer.constraint.AbstractRebalanceSoftConstraint;
 import org.apache.helix.api.rebalancer.constraint.dataprovider.PartitionWeightProvider;
 import org.apache.helix.common.ZkTestBase;
 import org.apache.helix.controller.common.PartitionStateMap;
 import org.apache.helix.controller.common.ResourcesStateMap;
 import org.apache.helix.controller.rebalancer.constraint.PartitionWeightAwareEvennessConstraint;
 import org.apache.helix.controller.rebalancer.constraint.TotalCapacityConstraint;
 import org.apache.helix.controller.rebalancer.constraint.dataprovider.MockCapacityProvider;
 import org.apache.helix.controller.rebalancer.constraint.dataprovider.MockPartitionWeightProvider;
 import org.apache.helix.controller.rebalancer.constraint.dataprovider.ZkBasedCapacityProvider;
 import org.apache.helix.controller.rebalancer.constraint.dataprovider.ZkBasedPartitionWeightProvider;
 import org.apache.helix.model.ClusterConfig;
 import org.apache.helix.model.InstanceConfig;
 import org.apache.helix.model.OnlineOfflineSMD;
 import org.apache.helix.model.Partition;
 import org.apache.helix.model.ResourceConfig;
 import org.apache.helix.util.WeightAwareRebalanceUtil;
 import org.testng.Assert;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;

 import static org.apache.helix.controller.rebalancer.constraint.dataprovider.ZkBasedPartitionWeightProvider.DEFAULT_WEIGHT_VALUE;

 public class TestWeightBasedRebalanceUtil extends ZkTestBase {
   private static String CLUSTER_NAME;

   final String resourceNamePrefix = "resource";
   final int nParticipants = 40;
   final int nResources = 20;
   final int nPartitions = 100;
   final int nReplicas = 3;
   final int defaultCapacity = 6000; // total = 6000*40 = 240000
   final int resourceWeight = 10; // total = 20*100*3*10 = 60000
   final String topState = "ONLINE";

   final List<String> resourceNames = new ArrayList<>();
   final List<String> instanceNames = new ArrayList<>();
   final List<String> partitions = new ArrayList<>(nPartitions);
   final List<ResourceConfig> resourceConfigs = new ArrayList<>();

   final LinkedHashMap<String, Integer> states = new LinkedHashMap<>(2);

   final ClusterConfig clusterConfig = new ClusterConfig(CLUSTER_NAME);
   final List<InstanceConfig> instanceConfigs = new ArrayList<>();

   @BeforeClass
   public void beforeClass() {
     System.out.println(
         "START " + getClass().getSimpleName() + " at " + new Date(System.currentTimeMillis()));

     CLUSTER_NAME = "MockCluster" + getShortClassName();

     for (int i = 0; i < nParticipants; i++) {
       instanceNames.add("node" + i);
     }
     for (int i = 0; i < nPartitions; i++) {
       partitions.add(Integer.toString(i));
     }

     for (int i = 0; i < nResources; i++) {
       resourceNames.add(resourceNamePrefix + i);
       ResourceConfig.Builder resourcBuilder = new ResourceConfig.Builder(resourceNamePrefix + i);
       resourcBuilder.setStateModelDefRef("OnlineOffline");
       resourcBuilder.setNumReplica(nReplicas);
       for (String partition : partitions) {
         resourcBuilder.setPreferenceList(partition, Collections.EMPTY_LIST);
       }
       resourceConfigs.add(resourcBuilder.build());
     }

     setupMockCluster();

     // setup storage cluster
     _gSetupTool.addCluster(CLUSTER_NAME, true);
   }

   @AfterClass
   public void afterClass() {
     deleteCluster(CLUSTER_NAME);
   }

   private void setupMockCluster() {
     for (String instance : instanceNames) {
       InstanceConfig config = new InstanceConfig(instance);
       instanceConfigs.add(config);
     }

     states.put("OFFLINE", 0);
     states.put(topState, nReplicas);
   }

   private Map<String, Integer> checkPartitionUsage(ResourcesStateMap assignment,
       PartitionWeightProvider weightProvider) {
     Map<String, Integer> weightCount = new HashMap<>();
     for (String resource : assignment.resourceSet()) {
       PartitionStateMap partitionMap = assignment.getPartitionStateMap(resource);
       for (Partition partition : partitionMap.partitionSet()) {
         // check states
         Map<String, Integer> stateCount = new HashMap<>(states);
         Map<String, String> stateMap = partitionMap.getPartitionMap(partition);
         for (String state : stateMap.values()) {
           Assert.assertTrue(stateCount.containsKey(state));
           stateCount.put(state, stateCount.get(state) - 1);
         }
         for (int count : stateCount.values()) {
           Assert.assertEquals(count, 0);
         }

         // report weight
         int partitionWeight =
             weightProvider.getPartitionWeight(resource, partition.getPartitionName());
         for (String instance : partitionMap.getPartitionMap(partition).keySet()) {
           if (!weightCount.containsKey(instance)) {
             weightCount.put(instance, partitionWeight);
           } else {
             weightCount.put(instance, weightCount.get(instance) + partitionWeight);
           }
         }
       }
     }
     return weightCount;
   }

   private void validateWeight(PartitionWeightProvider provider) {
     for (String resource : resourceNames) {
       for (String partition : partitions) {
         int weight = provider.getPartitionWeight(resource, partition);
         if (resource.equals(resourceNames.get(0))) {
           if (partition.equals(partitions.get(0))) {
             Assert.assertEquals(weight, resourceWeight * 3);
           } else {
             Assert.assertEquals(weight, resourceWeight * 2);
           }
         } else if (resource.equals(resourceNames.get(1))) {
           if (partition.equals(partitions.get(0))) {
             Assert.assertEquals(weight, resourceWeight * 3);
           } else {
             Assert.assertEquals(weight, resourceWeight);
           }
         } else {
           Assert.assertEquals(weight, resourceWeight);
         }
       }
     }
   }

   @Test
   public void testRebalance() {
     // capacity / weight
     Map<String, Integer> capacity = new HashMap<>();
     for (String instance : instanceNames) {
       capacity.put(instance, defaultCapacity);
     }

     MockCapacityProvider capacityProvider = new MockCapacityProvider(capacity, defaultCapacity);

     MockPartitionWeightProvider weightProvider = new MockPartitionWeightProvider(resourceWeight);

     TotalCapacityConstraint capacityConstraint =
         new TotalCapacityConstraint(weightProvider, capacityProvider);
     PartitionWeightAwareEvennessConstraint evenConstraint =
         new PartitionWeightAwareEvennessConstraint(weightProvider, capacityProvider);

     WeightAwareRebalanceUtil util = new WeightAwareRebalanceUtil(clusterConfig, instanceConfigs);
     ResourcesStateMap assignment = util.buildIncrementalRebalanceAssignment(resourceConfigs, null,
         Collections.<AbstractRebalanceHardConstraint> singletonList(capacityConstraint),
         Collections.<AbstractRebalanceSoftConstraint> singletonList(evenConstraint));
     Map<String, Integer> weightCount = checkPartitionUsage(assignment, weightProvider);

     int max = Collections.max(weightCount.values());
     int min = Collections.min(weightCount.values());
     // Since the accuracy of Default evenness constraint is 0.01, diff should be 1/100 of
     // participant capacity in max.
     Assert.assertTrue((max - min) <= defaultCapacity / 100);
   }

   @Test
   public void testZkBasedCapacityProvider() {
     Map<String, Integer> resourceDefaultWeightMap = new HashMap<>();
     resourceDefaultWeightMap.put(resourceNames.get(0), resourceWeight * 2);
     Map<String, Map<String, Integer>> partitionWeightMap = new HashMap<>();
     partitionWeightMap.put(resourceNames.get(0),
         Collections.singletonMap(partitions.get(0), resourceWeight * 3));
     partitionWeightMap.put(resourceNames.get(1),
         Collections.singletonMap(partitions.get(0), resourceWeight * 3));

     ZkBasedPartitionWeightProvider weightProvider =
         new ZkBasedPartitionWeightProvider(ZK_ADDR, CLUSTER_NAME, "Test");
     weightProvider.updateWeights(resourceDefaultWeightMap, partitionWeightMap, resourceWeight);
     // verify before persist
     validateWeight(weightProvider);

     // persist get values back
     weightProvider.persistWeights();
     // verify after persist
     weightProvider = new ZkBasedPartitionWeightProvider(ZK_ADDR, CLUSTER_NAME, "Test");
     validateWeight(weightProvider);

     weightProvider = new ZkBasedPartitionWeightProvider(ZK_ADDR, CLUSTER_NAME, "Fack");
     for (String resource : resourceNames) {
       for (String partition : partitions) {
         Assert.assertEquals(weightProvider.getPartitionWeight(resource, partition),
             DEFAULT_WEIGHT_VALUE);
       }
     }

     // update with invalid value
     weightProvider.updateWeights(Collections.EMPTY_MAP, Collections.EMPTY_MAP, -1);
     try {
       weightProvider.persistWeights();
       Assert.fail("Should fail to persist invalid weight information.");
     } catch (HelixException hex) {
       // expected
     }

     Map<String, Integer> capacity = new HashMap<>();
     Map<String, Integer> usage = new HashMap<>();
     for (int i = 0; i < instanceNames.size(); i++) {
       capacity.put(instanceNames.get(i), defaultCapacity + i);
       usage.put(instanceNames.get(i), i);
     }

     ZkBasedCapacityProvider capacityProvider =
         new ZkBasedCapacityProvider(ZK_ADDR, CLUSTER_NAME, "Test");
     capacityProvider.updateCapacity(capacity, usage, defaultCapacity);

     for (String instance : instanceNames) {
       Assert.assertEquals(capacityProvider.getParticipantCapacity(instance),
           capacity.get(instance).intValue());
       Assert.assertEquals(capacityProvider.getParticipantUsage(instance),
           usage.get(instance).intValue());
     }

     // persist get values back
     capacityProvider.persistCapacity();
     capacityProvider = new ZkBasedCapacityProvider(ZK_ADDR, CLUSTER_NAME, "Test");
     for (String instance : instanceNames) {
       Assert.assertEquals(capacityProvider.getParticipantCapacity(instance),
           capacity.get(instance).intValue());
       Assert.assertEquals(capacityProvider.getParticipantUsage(instance),
           usage.get(instance).intValue());
     }

     // update usage
     String targetInstanceName = instanceNames.get(0);
     int newUsgae = 12345;
     capacityProvider.updateCapacity(Collections.EMPTY_MAP,
         Collections.singletonMap(targetInstanceName, newUsgae), defaultCapacity);
     Assert.assertEquals(capacityProvider.getParticipantUsage(targetInstanceName), newUsgae);
     // check again without updating ZK
     capacityProvider = new ZkBasedCapacityProvider(ZK_ADDR, CLUSTER_NAME, "Test");
     Assert.assertEquals(capacityProvider.getParticipantUsage(targetInstanceName), 0);

     // update with invalid value
     capacityProvider.updateCapacity(Collections.EMPTY_MAP, Collections.EMPTY_MAP, -1);
     try {
       capacityProvider.persistCapacity();
       Assert.fail("Should fail to persist invalid weight information.");
     } catch (HelixException hex) {
       // expected
     }
   }

   @Test
   public void testRebalanceUsingZkDataProvider() {
     // capacity / weight
     Map<String, Integer> capacity = new HashMap<>();
     for (String instance : instanceNames) {
       capacity.put(instance, defaultCapacity);
     }

     ZkBasedPartitionWeightProvider weightProvider =
         new ZkBasedPartitionWeightProvider(ZK_ADDR, CLUSTER_NAME, "QPS");
     weightProvider.updateWeights(Collections.EMPTY_MAP, Collections.EMPTY_MAP, resourceWeight);

     ZkBasedCapacityProvider capacityProvider =
         new ZkBasedCapacityProvider(ZK_ADDR, CLUSTER_NAME, "QPS");
     capacityProvider.updateCapacity(capacity, Collections.EMPTY_MAP, 0);

     TotalCapacityConstraint capacityConstraint =
         new TotalCapacityConstraint(weightProvider, capacityProvider);
     PartitionWeightAwareEvennessConstraint evenConstraint =
         new PartitionWeightAwareEvennessConstraint(weightProvider, capacityProvider);

     WeightAwareRebalanceUtil util = new WeightAwareRebalanceUtil(clusterConfig, instanceConfigs);
     ResourcesStateMap assignment = util.buildIncrementalRebalanceAssignment(resourceConfigs, null,
         Collections.<AbstractRebalanceHardConstraint> singletonList(capacityConstraint),
         Collections.<AbstractRebalanceSoftConstraint> singletonList(evenConstraint));
     Map<String, Integer> weightCount = checkPartitionUsage(assignment, weightProvider);

     int max = Collections.max(weightCount.values());
     int min = Collections.min(weightCount.values());
     // Since the accuracy of Default evenness constraint is 0.01, diff should be 1/100 of
     // participant capacity in max.
     Assert.assertTrue((max - min) <= defaultCapacity / 100);
   }

   @Test(dependsOnMethods = "testRebalanceUsingZkDataProvider")
   public void testRebalanceWithExistingUsage() {
     // capacity / weight
     Map<String, Integer> capacity = new HashMap<>();
     Map<String, Integer> usage = new HashMap<>();
     for (int i = 0; i < instanceNames.size(); i++) {
       String instance = instanceNames.get(i);
       capacity.put(instance, defaultCapacity);
       if (i % 7 == 0) {
         usage.put(instance, defaultCapacity);
       }
     }

     ZkBasedPartitionWeightProvider weightProvider =
         new ZkBasedPartitionWeightProvider(ZK_ADDR, CLUSTER_NAME, "QPS");
     weightProvider.updateWeights(Collections.EMPTY_MAP, Collections.EMPTY_MAP, resourceWeight);

     ZkBasedCapacityProvider capacityProvider =
         new ZkBasedCapacityProvider(ZK_ADDR, CLUSTER_NAME, "QPS");
     capacityProvider.updateCapacity(capacity, usage, 0);

     TotalCapacityConstraint hardConstraint =
         new TotalCapacityConstraint(weightProvider, capacityProvider);
     PartitionWeightAwareEvennessConstraint evenConstraint =
         new PartitionWeightAwareEvennessConstraint(weightProvider, capacityProvider);

     WeightAwareRebalanceUtil util = new WeightAwareRebalanceUtil(clusterConfig, instanceConfigs);
     ResourcesStateMap assignment = util.buildIncrementalRebalanceAssignment(resourceConfigs, null,
         Collections.<AbstractRebalanceHardConstraint> singletonList(hardConstraint),
         Collections.<AbstractRebalanceSoftConstraint> singletonList(evenConstraint));
     Map<String, Integer> weightCount = checkPartitionUsage(assignment, weightProvider);

     for (int i = 0; i < instanceNames.size(); i++) {
       String instance = instanceNames.get(i);
       if (i % 7 == 0) {
         Assert.assertFalse(weightCount.containsKey(instance));
       } else {
         Assert.assertTrue(weightCount.get(instance) > 0);
       }
     }
   }

   @Test(dependsOnMethods = "testRebalanceUsingZkDataProvider")
   public void testRebalanceOption() {
     // capacity / weight
     Map<String, Integer> capacity = new HashMap<>();
     for (String instance : instanceNames) {
       capacity.put(instance, defaultCapacity);
     }

     ZkBasedPartitionWeightProvider weightProvider =
         new ZkBasedPartitionWeightProvider(ZK_ADDR, CLUSTER_NAME, "QPS");
     weightProvider.updateWeights(Collections.EMPTY_MAP, Collections.EMPTY_MAP, resourceWeight);

     ZkBasedCapacityProvider capacityProvider =
         new ZkBasedCapacityProvider(ZK_ADDR, CLUSTER_NAME, "QPS");
     capacityProvider.updateCapacity(capacity, Collections.EMPTY_MAP, 0);

     PartitionWeightAwareEvennessConstraint evenConstraint =
         new PartitionWeightAwareEvennessConstraint(weightProvider, capacityProvider);

     // Assume existing assignment
     ResourcesStateMap existingAssignment = new ResourcesStateMap();
     String targetResource = resourceNames.get(0);
     for (String partition : partitions) {
       for (int i = 0; i < nReplicas; i++) {
         existingAssignment.setState(targetResource, new Partition(partition), instanceNames.get(i),
             topState);
       }
     }

     WeightAwareRebalanceUtil util = new WeightAwareRebalanceUtil(clusterConfig, instanceConfigs);

     // INCREMENTAL
     ResourcesStateMap assignment = util.buildIncrementalRebalanceAssignment(resourceConfigs,
         existingAssignment, Collections.EMPTY_LIST,
         Collections.<AbstractRebalanceSoftConstraint> singletonList(evenConstraint));
     // check if the existingAssignment is changed
     for (String partition : partitions) {
       Assert.assertTrue(assignment.getInstanceStateMap(targetResource, new Partition(partition))
           .keySet().containsAll(instanceNames.subList(0, nReplicas)));
     }
     // still need to check for balance
     Map<String, Integer> weightCount = checkPartitionUsage(assignment, weightProvider);
     int max = Collections.max(weightCount.values());
     int min = Collections.min(weightCount.values());
     // Since the accuracy of Default evenness constraint is 0.01, diff should be 1/100 of
     // participant capacity in max.
     Assert.assertTrue((max - min) <= defaultCapacity / 100);

     // FULL
     assignment = util.buildFullRebalanceAssignment(resourceConfigs, existingAssignment,
         Collections.EMPTY_LIST,
         Collections.<AbstractRebalanceSoftConstraint> singletonList(evenConstraint));
     // check if the existingAssignment is changed
     for (String partition : partitions) {
       Assert.assertFalse(assignment.getInstanceStateMap(targetResource, new Partition(partition))
           .keySet().containsAll(instanceNames.subList(0, nReplicas)));
     }
   }

   @Test(dependsOnMethods = "testRebalanceUsingZkDataProvider")
   public void testInvalidInput() {
     // capacity / weight
     Map<String, Integer> capacity = new HashMap<>();
     for (String instance : instanceNames) {
       capacity.put(instance, defaultCapacity);
     }

     ZkBasedPartitionWeightProvider weightProvider =
         new ZkBasedPartitionWeightProvider(ZK_ADDR, CLUSTER_NAME, "QPS");
     weightProvider.updateWeights(Collections.EMPTY_MAP, Collections.EMPTY_MAP, resourceWeight);

     ZkBasedCapacityProvider capacityProvider =
         new ZkBasedCapacityProvider(ZK_ADDR, CLUSTER_NAME, "QPS");
     capacityProvider.updateCapacity(capacity, Collections.EMPTY_MAP, 0);

     TotalCapacityConstraint capacityConstraint =
         new TotalCapacityConstraint(weightProvider, capacityProvider);

     WeightAwareRebalanceUtil util = new WeightAwareRebalanceUtil(clusterConfig, instanceConfigs);

     // Empty constraint
     try {
       util.buildIncrementalRebalanceAssignment(resourceConfigs, null, Collections.EMPTY_LIST,
           Collections.EMPTY_LIST);
       Assert.fail("Should fail due to empty constraint list.");
     } catch (HelixException ex) {
       // expected
     }

     ResourceConfig.Builder invalidResourceBuilder = new ResourceConfig.Builder("InvalidResource");
     invalidResourceBuilder.setStateModelDefRef("OnlineOffline");
     invalidResourceBuilder.setNumPartitions(nPartitions);
     invalidResourceBuilder.setNumReplica(nReplicas);
     for (String partition : partitions) {
       invalidResourceBuilder.setPreferenceList(partition, Collections.EMPTY_LIST);
     }

     // Auto mode resource config
     try {
       invalidResourceBuilder
           .setRebalanceConfig(new RebalanceConfig(new ZNRecord("InvalidResource")));
       invalidResourceBuilder.getRebalanceConfig()
           .setRebalanceMode(RebalanceConfig.RebalanceMode.FULL_AUTO);
       util.buildIncrementalRebalanceAssignment(
           Collections.singletonList(invalidResourceBuilder.build()), null,
           Collections.<AbstractRebalanceHardConstraint> singletonList(capacityConstraint),
           Collections.EMPTY_LIST);
       Assert.fail("Should fail due to full auto resource config.");
     } catch (HelixException ex) {
       // expected
       invalidResourceBuilder.getRebalanceConfig()
           .setRebalanceMode(RebalanceConfig.RebalanceMode.CUSTOMIZED);
     }

     // Auto mode resource config
     try {
       invalidResourceBuilder.setStateModelDefRef("CustomizedOnlineOffline");
       util.buildIncrementalRebalanceAssignment(
           Collections.singletonList(invalidResourceBuilder.build()), null,
           Collections.<AbstractRebalanceHardConstraint> singletonList(capacityConstraint),
           Collections.EMPTY_LIST);
       Assert.fail("Should fail due to unknown state model def ref.");
     } catch (IllegalArgumentException ex) {
       // expected
       util.registerCustomizedStateModelDef("CustomizedOnlineOffline", OnlineOfflineSMD.build());
       ResourcesStateMap assignment = util.buildIncrementalRebalanceAssignment(
           Collections.singletonList(invalidResourceBuilder.build()), null,
           Collections.<AbstractRebalanceHardConstraint> singletonList(capacityConstraint),
           Collections.EMPTY_LIST);
       checkPartitionUsage(assignment, weightProvider);
     }
   }
 }
	package org.apache.helix.integration;

	/*
	* 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.
	*/

	import java.util.ArrayList;
	import java.util.Collections;
	import java.util.Date;
	import java.util.HashMap;
	import java.util.LinkedHashMap;
	import java.util.List;
	import java.util.Map;

	import org.apache.helix.HelixException;
	import org.apache.helix.zookeeper.datamodel.ZNRecord;
	import org.apache.helix.api.config.RebalanceConfig;
	import org.apache.helix.api.rebalancer.constraint.AbstractRebalanceHardConstraint;
	import org.apache.helix.api.rebalancer.constraint.AbstractRebalanceSoftConstraint;
	import org.apache.helix.api.rebalancer.constraint.dataprovider.PartitionWeightProvider;
	import org.apache.helix.common.ZkTestBase;
	import org.apache.helix.controller.common.PartitionStateMap;
	import org.apache.helix.controller.common.ResourcesStateMap;
	import org.apache.helix.controller.rebalancer.constraint.PartitionWeightAwareEvennessConstraint;
	import org.apache.helix.controller.rebalancer.constraint.TotalCapacityConstraint;
	import org.apache.helix.controller.rebalancer.constraint.dataprovider.MockCapacityProvider;
	import org.apache.helix.controller.rebalancer.constraint.dataprovider.MockPartitionWeightProvider;
	import org.apache.helix.controller.rebalancer.constraint.dataprovider.ZkBasedCapacityProvider;
	import org.apache.helix.controller.rebalancer.constraint.dataprovider.ZkBasedPartitionWeightProvider;
	import org.apache.helix.model.ClusterConfig;
	import org.apache.helix.model.InstanceConfig;
	import org.apache.helix.model.OnlineOfflineSMD;
	import org.apache.helix.model.Partition;
	import org.apache.helix.model.ResourceConfig;
	import org.apache.helix.util.WeightAwareRebalanceUtil;
	import org.testng.Assert;
	import org.testng.annotations.AfterClass;
	import org.testng.annotations.BeforeClass;
	import org.testng.annotations.Test;

	import static org.apache.helix.controller.rebalancer.constraint.dataprovider.ZkBasedPartitionWeightProvider.DEFAULT_WEIGHT_VALUE;

	public class TestWeightBasedRebalanceUtil extends ZkTestBase {
	private static String CLUSTER_NAME;

	final String resourceNamePrefix = "resource";
	final int nParticipants = 40;
	final int nResources = 20;
	final int nPartitions = 100;
	final int nReplicas = 3;
	final int defaultCapacity = 6000; // total = 6000*40 = 240000
	final int resourceWeight = 10; // total = 201003*10 = 60000
	final String topState = "ONLINE";

	final List<String> resourceNames = new ArrayList<>();
	final List<String> instanceNames = new ArrayList<>();
	final List<String> partitions = new ArrayList<>(nPartitions);
	final List<ResourceConfig> resourceConfigs = new ArrayList<>();

	final LinkedHashMap<String, Integer> states = new LinkedHashMap<>(2);

	final ClusterConfig clusterConfig = new ClusterConfig(CLUSTER_NAME);
	final List<InstanceConfig> instanceConfigs = new ArrayList<>();

	@BeforeClass
	public void beforeClass() {
	System.out.println(
	"START " + getClass().getSimpleName() + " at " + new Date(System.currentTimeMillis()));

	CLUSTER_NAME = "MockCluster" + getShortClassName();

	for (int i = 0; i < nParticipants; i++) {
	instanceNames.add("node" + i);
	}
	for (int i = 0; i < nPartitions; i++) {
	partitions.add(Integer.toString(i));
	}

	for (int i = 0; i < nResources; i++) {
	resourceNames.add(resourceNamePrefix + i);
	ResourceConfig.Builder resourcBuilder = new ResourceConfig.Builder(resourceNamePrefix + i);
	resourcBuilder.setStateModelDefRef("OnlineOffline");
	resourcBuilder.setNumReplica(nReplicas);
	for (String partition : partitions) {
	resourcBuilder.setPreferenceList(partition, Collections.EMPTY_LIST);
	}
	resourceConfigs.add(resourcBuilder.build());
	}

	setupMockCluster();

	// setup storage cluster
	_gSetupTool.addCluster(CLUSTER_NAME, true);
	}

	@AfterClass
	public void afterClass() {
	deleteCluster(CLUSTER_NAME);
	}

	private void setupMockCluster() {
	for (String instance : instanceNames) {
	InstanceConfig config = new InstanceConfig(instance);
	instanceConfigs.add(config);
	}

	states.put("OFFLINE", 0);
	states.put(topState, nReplicas);
	}

	private Map<String, Integer> checkPartitionUsage(ResourcesStateMap assignment,
	PartitionWeightProvider weightProvider) {
	Map<String, Integer> weightCount = new HashMap<>();
	for (String resource : assignment.resourceSet()) {
	PartitionStateMap partitionMap = assignment.getPartitionStateMap(resource);
	for (Partition partition : partitionMap.partitionSet()) {
	// check states
	Map<String, Integer> stateCount = new HashMap<>(states);
	Map<String, String> stateMap = partitionMap.getPartitionMap(partition);
	for (String state : stateMap.values()) {
	Assert.assertTrue(stateCount.containsKey(state));
	stateCount.put(state, stateCount.get(state) - 1);
	}
	for (int count : stateCount.values()) {
	Assert.assertEquals(count, 0);
	}

	// report weight
	int partitionWeight =
	weightProvider.getPartitionWeight(resource, partition.getPartitionName());
	for (String instance : partitionMap.getPartitionMap(partition).keySet()) {
	if (!weightCount.containsKey(instance)) {
	weightCount.put(instance, partitionWeight);
	} else {
	weightCount.put(instance, weightCount.get(instance) + partitionWeight);
	}
	}
	}
	}
	return weightCount;
	}

	private void validateWeight(PartitionWeightProvider provider) {
	for (String resource : resourceNames) {
	for (String partition : partitions) {
	int weight = provider.getPartitionWeight(resource, partition);
	if (resource.equals(resourceNames.get(0))) {
	if (partition.equals(partitions.get(0))) {
	Assert.assertEquals(weight, resourceWeight * 3);
	} else {
	Assert.assertEquals(weight, resourceWeight * 2);
	}
	} else if (resource.equals(resourceNames.get(1))) {
	if (partition.equals(partitions.get(0))) {
	Assert.assertEquals(weight, resourceWeight * 3);
	} else {
	Assert.assertEquals(weight, resourceWeight);
	}
	} else {
	Assert.assertEquals(weight, resourceWeight);
	}
	}
	}
	}

	@Test
	public void testRebalance() {
	// capacity / weight
	Map<String, Integer> capacity = new HashMap<>();
	for (String instance : instanceNames) {
	capacity.put(instance, defaultCapacity);
	}

	MockCapacityProvider capacityProvider = new MockCapacityProvider(capacity, defaultCapacity);

	MockPartitionWeightProvider weightProvider = new MockPartitionWeightProvider(resourceWeight);

	TotalCapacityConstraint capacityConstraint =
	new TotalCapacityConstraint(weightProvider, capacityProvider);
	PartitionWeightAwareEvennessConstraint evenConstraint =
	new PartitionWeightAwareEvennessConstraint(weightProvider, capacityProvider);

	WeightAwareRebalanceUtil util = new WeightAwareRebalanceUtil(clusterConfig, instanceConfigs);
	ResourcesStateMap assignment = util.buildIncrementalRebalanceAssignment(resourceConfigs, null,
	Collections.<AbstractRebalanceHardConstraint> singletonList(capacityConstraint),
	Collections.<AbstractRebalanceSoftConstraint> singletonList(evenConstraint));
	Map<String, Integer> weightCount = checkPartitionUsage(assignment, weightProvider);

	int max = Collections.max(weightCount.values());
	int min = Collections.min(weightCount.values());
	// Since the accuracy of Default evenness constraint is 0.01, diff should be 1/100 of
	// participant capacity in max.
	Assert.assertTrue((max - min) <= defaultCapacity / 100);
	}

	@Test
	public void testZkBasedCapacityProvider() {
	Map<String, Integer> resourceDefaultWeightMap = new HashMap<>();
	resourceDefaultWeightMap.put(resourceNames.get(0), resourceWeight * 2);
	Map<String, Map<String, Integer>> partitionWeightMap = new HashMap<>();
	partitionWeightMap.put(resourceNames.get(0),
	Collections.singletonMap(partitions.get(0), resourceWeight * 3));
	partitionWeightMap.put(resourceNames.get(1),
	Collections.singletonMap(partitions.get(0), resourceWeight * 3));

	ZkBasedPartitionWeightProvider weightProvider =
	new ZkBasedPartitionWeightProvider(ZK_ADDR, CLUSTER_NAME, "Test");
	weightProvider.updateWeights(resourceDefaultWeightMap, partitionWeightMap, resourceWeight);
	// verify before persist
	validateWeight(weightProvider);

	// persist get values back
	weightProvider.persistWeights();
	// verify after persist
	weightProvider = new ZkBasedPartitionWeightProvider(ZK_ADDR, CLUSTER_NAME, "Test");
	validateWeight(weightProvider);

	weightProvider = new ZkBasedPartitionWeightProvider(ZK_ADDR, CLUSTER_NAME, "Fack");
	for (String resource : resourceNames) {
	for (String partition : partitions) {
	Assert.assertEquals(weightProvider.getPartitionWeight(resource, partition),
	DEFAULT_WEIGHT_VALUE);
	}
	}

	// update with invalid value
	weightProvider.updateWeights(Collections.EMPTY_MAP, Collections.EMPTY_MAP, -1);
	try {
	weightProvider.persistWeights();
	Assert.fail("Should fail to persist invalid weight information.");
	} catch (HelixException hex) {
	// expected
	}

	Map<String, Integer> capacity = new HashMap<>();
	Map<String, Integer> usage = new HashMap<>();
	for (int i = 0; i < instanceNames.size(); i++) {
	capacity.put(instanceNames.get(i), defaultCapacity + i);
	usage.put(instanceNames.get(i), i);
	}

	ZkBasedCapacityProvider capacityProvider =
	new ZkBasedCapacityProvider(ZK_ADDR, CLUSTER_NAME, "Test");
	capacityProvider.updateCapacity(capacity, usage, defaultCapacity);

	for (String instance : instanceNames) {
	Assert.assertEquals(capacityProvider.getParticipantCapacity(instance),
	capacity.get(instance).intValue());
	Assert.assertEquals(capacityProvider.getParticipantUsage(instance),
	usage.get(instance).intValue());
	}

	// persist get values back
	capacityProvider.persistCapacity();
	capacityProvider = new ZkBasedCapacityProvider(ZK_ADDR, CLUSTER_NAME, "Test");
	for (String instance : instanceNames) {
	Assert.assertEquals(capacityProvider.getParticipantCapacity(instance),
	capacity.get(instance).intValue());
	Assert.assertEquals(capacityProvider.getParticipantUsage(instance),
	usage.get(instance).intValue());
	}

	// update usage
	String targetInstanceName = instanceNames.get(0);
	int newUsgae = 12345;
	capacityProvider.updateCapacity(Collections.EMPTY_MAP,
	Collections.singletonMap(targetInstanceName, newUsgae), defaultCapacity);
	Assert.assertEquals(capacityProvider.getParticipantUsage(targetInstanceName), newUsgae);
	// check again without updating ZK
	capacityProvider = new ZkBasedCapacityProvider(ZK_ADDR, CLUSTER_NAME, "Test");
	Assert.assertEquals(capacityProvider.getParticipantUsage(targetInstanceName), 0);

	// update with invalid value
	capacityProvider.updateCapacity(Collections.EMPTY_MAP, Collections.EMPTY_MAP, -1);
	try {
	capacityProvider.persistCapacity();
	Assert.fail("Should fail to persist invalid weight information.");
	} catch (HelixException hex) {
	// expected
	}
	}

	@Test
	public void testRebalanceUsingZkDataProvider() {
	// capacity / weight
	Map<String, Integer> capacity = new HashMap<>();
	for (String instance : instanceNames) {
	capacity.put(instance, defaultCapacity);
	}

	ZkBasedPartitionWeightProvider weightProvider =
	new ZkBasedPartitionWeightProvider(ZK_ADDR, CLUSTER_NAME, "QPS");
	weightProvider.updateWeights(Collections.EMPTY_MAP, Collections.EMPTY_MAP, resourceWeight);

	ZkBasedCapacityProvider capacityProvider =
	new ZkBasedCapacityProvider(ZK_ADDR, CLUSTER_NAME, "QPS");
	capacityProvider.updateCapacity(capacity, Collections.EMPTY_MAP, 0);

	TotalCapacityConstraint capacityConstraint =
	new TotalCapacityConstraint(weightProvider, capacityProvider);
	PartitionWeightAwareEvennessConstraint evenConstraint =
	new PartitionWeightAwareEvennessConstraint(weightProvider, capacityProvider);

	WeightAwareRebalanceUtil util = new WeightAwareRebalanceUtil(clusterConfig, instanceConfigs);
	ResourcesStateMap assignment = util.buildIncrementalRebalanceAssignment(resourceConfigs, null,
	Collections.<AbstractRebalanceHardConstraint> singletonList(capacityConstraint),
	Collections.<AbstractRebalanceSoftConstraint> singletonList(evenConstraint));
	Map<String, Integer> weightCount = checkPartitionUsage(assignment, weightProvider);

	int max = Collections.max(weightCount.values());
	int min = Collections.min(weightCount.values());
	// Since the accuracy of Default evenness constraint is 0.01, diff should be 1/100 of
	// participant capacity in max.
	Assert.assertTrue((max - min) <= defaultCapacity / 100);
	}

	@Test(dependsOnMethods = "testRebalanceUsingZkDataProvider")
	public void testRebalanceWithExistingUsage() {
	// capacity / weight
	Map<String, Integer> capacity = new HashMap<>();
	Map<String, Integer> usage = new HashMap<>();
	for (int i = 0; i < instanceNames.size(); i++) {
	String instance = instanceNames.get(i);
	capacity.put(instance, defaultCapacity);
	if (i % 7 == 0) {
	usage.put(instance, defaultCapacity);
	}
	}

	ZkBasedPartitionWeightProvider weightProvider =
	new ZkBasedPartitionWeightProvider(ZK_ADDR, CLUSTER_NAME, "QPS");
	weightProvider.updateWeights(Collections.EMPTY_MAP, Collections.EMPTY_MAP, resourceWeight);

	ZkBasedCapacityProvider capacityProvider =
	new ZkBasedCapacityProvider(ZK_ADDR, CLUSTER_NAME, "QPS");
	capacityProvider.updateCapacity(capacity, usage, 0);

	TotalCapacityConstraint hardConstraint =
	new TotalCapacityConstraint(weightProvider, capacityProvider);
	PartitionWeightAwareEvennessConstraint evenConstraint =
	new PartitionWeightAwareEvennessConstraint(weightProvider, capacityProvider);

	WeightAwareRebalanceUtil util = new WeightAwareRebalanceUtil(clusterConfig, instanceConfigs);
	ResourcesStateMap assignment = util.buildIncrementalRebalanceAssignment(resourceConfigs, null,
	Collections.<AbstractRebalanceHardConstraint> singletonList(hardConstraint),
	Collections.<AbstractRebalanceSoftConstraint> singletonList(evenConstraint));
	Map<String, Integer> weightCount = checkPartitionUsage(assignment, weightProvider);

	for (int i = 0; i < instanceNames.size(); i++) {
	String instance = instanceNames.get(i);
	if (i % 7 == 0) {
	Assert.assertFalse(weightCount.containsKey(instance));
	} else {
	Assert.assertTrue(weightCount.get(instance) > 0);
	}
	}
	}

	@Test(dependsOnMethods = "testRebalanceUsingZkDataProvider")
	public void testRebalanceOption() {
	// capacity / weight
	Map<String, Integer> capacity = new HashMap<>();
	for (String instance : instanceNames) {
	capacity.put(instance, defaultCapacity);
	}

	ZkBasedPartitionWeightProvider weightProvider =
	new ZkBasedPartitionWeightProvider(ZK_ADDR, CLUSTER_NAME, "QPS");
	weightProvider.updateWeights(Collections.EMPTY_MAP, Collections.EMPTY_MAP, resourceWeight);

	ZkBasedCapacityProvider capacityProvider =
	new ZkBasedCapacityProvider(ZK_ADDR, CLUSTER_NAME, "QPS");
	capacityProvider.updateCapacity(capacity, Collections.EMPTY_MAP, 0);

	PartitionWeightAwareEvennessConstraint evenConstraint =
	new PartitionWeightAwareEvennessConstraint(weightProvider, capacityProvider);

	// Assume existing assignment
	ResourcesStateMap existingAssignment = new ResourcesStateMap();
	String targetResource = resourceNames.get(0);
	for (String partition : partitions) {
	for (int i = 0; i < nReplicas; i++) {
	existingAssignment.setState(targetResource, new Partition(partition), instanceNames.get(i),
	topState);
	}
	}

	WeightAwareRebalanceUtil util = new WeightAwareRebalanceUtil(clusterConfig, instanceConfigs);

	// INCREMENTAL
	ResourcesStateMap assignment = util.buildIncrementalRebalanceAssignment(resourceConfigs,
	existingAssignment, Collections.EMPTY_LIST,
	Collections.<AbstractRebalanceSoftConstraint> singletonList(evenConstraint));
	// check if the existingAssignment is changed
	for (String partition : partitions) {
	Assert.assertTrue(assignment.getInstanceStateMap(targetResource, new Partition(partition))
	.keySet().containsAll(instanceNames.subList(0, nReplicas)));
	}
	// still need to check for balance
	Map<String, Integer> weightCount = checkPartitionUsage(assignment, weightProvider);
	int max = Collections.max(weightCount.values());
	int min = Collections.min(weightCount.values());
	// Since the accuracy of Default evenness constraint is 0.01, diff should be 1/100 of
	// participant capacity in max.
	Assert.assertTrue((max - min) <= defaultCapacity / 100);

	// FULL
	assignment = util.buildFullRebalanceAssignment(resourceConfigs, existingAssignment,
	Collections.EMPTY_LIST,
	Collections.<AbstractRebalanceSoftConstraint> singletonList(evenConstraint));
	// check if the existingAssignment is changed
	for (String partition : partitions) {
	Assert.assertFalse(assignment.getInstanceStateMap(targetResource, new Partition(partition))
	.keySet().containsAll(instanceNames.subList(0, nReplicas)));
	}
	}

	@Test(dependsOnMethods = "testRebalanceUsingZkDataProvider")
	public void testInvalidInput() {
	// capacity / weight
	Map<String, Integer> capacity = new HashMap<>();
	for (String instance : instanceNames) {
	capacity.put(instance, defaultCapacity);
	}

	ZkBasedPartitionWeightProvider weightProvider =
	new ZkBasedPartitionWeightProvider(ZK_ADDR, CLUSTER_NAME, "QPS");
	weightProvider.updateWeights(Collections.EMPTY_MAP, Collections.EMPTY_MAP, resourceWeight);

	ZkBasedCapacityProvider capacityProvider =
	new ZkBasedCapacityProvider(ZK_ADDR, CLUSTER_NAME, "QPS");
	capacityProvider.updateCapacity(capacity, Collections.EMPTY_MAP, 0);

	TotalCapacityConstraint capacityConstraint =
	new TotalCapacityConstraint(weightProvider, capacityProvider);

	WeightAwareRebalanceUtil util = new WeightAwareRebalanceUtil(clusterConfig, instanceConfigs);

	// Empty constraint
	try {
	util.buildIncrementalRebalanceAssignment(resourceConfigs, null, Collections.EMPTY_LIST,
	Collections.EMPTY_LIST);
	Assert.fail("Should fail due to empty constraint list.");
	} catch (HelixException ex) {
	// expected
	}

	ResourceConfig.Builder invalidResourceBuilder = new ResourceConfig.Builder("InvalidResource");
	invalidResourceBuilder.setStateModelDefRef("OnlineOffline");
	invalidResourceBuilder.setNumPartitions(nPartitions);
	invalidResourceBuilder.setNumReplica(nReplicas);
	for (String partition : partitions) {
	invalidResourceBuilder.setPreferenceList(partition, Collections.EMPTY_LIST);
	}

	// Auto mode resource config
	try {
	invalidResourceBuilder
	.setRebalanceConfig(new RebalanceConfig(new ZNRecord("InvalidResource")));
	invalidResourceBuilder.getRebalanceConfig()
	.setRebalanceMode(RebalanceConfig.RebalanceMode.FULL_AUTO);
	util.buildIncrementalRebalanceAssignment(
	Collections.singletonList(invalidResourceBuilder.build()), null,
	Collections.<AbstractRebalanceHardConstraint> singletonList(capacityConstraint),
	Collections.EMPTY_LIST);
	Assert.fail("Should fail due to full auto resource config.");
	} catch (HelixException ex) {
	// expected
	invalidResourceBuilder.getRebalanceConfig()
	.setRebalanceMode(RebalanceConfig.RebalanceMode.CUSTOMIZED);
	}

	// Auto mode resource config
	try {
	invalidResourceBuilder.setStateModelDefRef("CustomizedOnlineOffline");
	util.buildIncrementalRebalanceAssignment(
	Collections.singletonList(invalidResourceBuilder.build()), null,
	Collections.<AbstractRebalanceHardConstraint> singletonList(capacityConstraint),
	Collections.EMPTY_LIST);
	Assert.fail("Should fail due to unknown state model def ref.");
	} catch (IllegalArgumentException ex) {
	// expected
	util.registerCustomizedStateModelDef("CustomizedOnlineOffline", OnlineOfflineSMD.build());
	ResourcesStateMap assignment = util.buildIncrementalRebalanceAssignment(
	Collections.singletonList(invalidResourceBuilder.build()), null,
	Collections.<AbstractRebalanceHardConstraint> singletonList(capacityConstraint),
	Collections.EMPTY_LIST);
	checkPartitionUsage(assignment, weightProvider);
	}
	}
	}