helix-core/src/test/java/org/apache/helix/integration/TestPreferenceListAsQueue.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.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Queue;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;

 import com.google.common.collect.Lists;
 import org.apache.helix.AccessOption;
 import org.apache.helix.HelixAdmin;
 import org.apache.helix.HelixDataAccessor;
 import org.apache.helix.HelixManager;
 import org.apache.helix.HelixManagerFactory;
 import org.apache.helix.InstanceType;
 import org.apache.helix.NotificationContext;
 import org.apache.helix.PropertyKey;
 import org.apache.helix.TestHelper;
 import org.apache.helix.zookeeper.datamodel.ZNRecord;
 import org.apache.helix.ZkUnitTestBase;
 import org.apache.helix.model.ClusterConstraints.ConstraintAttribute;
 import org.apache.helix.model.ClusterConstraints.ConstraintType;
 import org.apache.helix.model.IdealState;
 import org.apache.helix.model.IdealState.IdealStateProperty;
 import org.apache.helix.model.IdealState.RebalanceMode;
 import org.apache.helix.model.Message;
 import org.apache.helix.model.StateModelDefinition;
 import org.apache.helix.model.builder.ConstraintItemBuilder;
 import org.apache.helix.participant.statemachine.StateModel;
 import org.apache.helix.participant.statemachine.StateModelFactory;
 import org.apache.helix.tools.ClusterSetup;
 import org.apache.helix.zookeeper.zkclient.DataUpdater;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.testng.Assert;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;

 public class TestPreferenceListAsQueue extends ZkUnitTestBase {
   private static final Logger LOG = LoggerFactory.getLogger(TestPreferenceListAsQueue.class);
   private static final int TRANSITION_TIME = 500;
   private static final int PARALLELISM = 1;

   private List<String> _instanceList;
   private Queue<List<String>> _prefListHistory;
   private String _clusterName;
   private String _stateModel;
   private ClusterSetup _clusterSetup;
   private HelixAdmin _admin;
   private CountDownLatch _onlineLatch;
   private CountDownLatch _offlineLatch;

   @BeforeMethod
   public void beforeMethod() {
     _instanceList = Lists.newLinkedList();
     _clusterSetup = new ClusterSetup(ZK_ADDR);
     _admin = _clusterSetup.getClusterManagementTool();
     _prefListHistory = Lists.newLinkedList();

     // Create cluster
     String className = TestHelper.getTestClassName();
     String methodName = TestHelper.getTestMethodName();
     _clusterName = className + "_" + methodName;
     _clusterSetup.addCluster(_clusterName, true);
   }

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

   /**
    * This test ensures successful creation when the state model has OFFLINE --> deprioritized and
    * a partition-level constraint enforces parallelism
    * @throws Exception
    */
   @Test
   public void testReprioritizedWithConstraint() throws Exception {
     _stateModel = "OnlineOfflineReprioritized";

     // Add a state model with the transition to ONLINE deprioritized
     _clusterSetup.addStateModelDef(_clusterName, _stateModel,
         createReprioritizedStateModelDef(_stateModel));

     // Add a constraint of one transition per partition
     ConstraintItemBuilder constraintItemBuilder = new ConstraintItemBuilder();
     constraintItemBuilder
         .addConstraintAttribute(ConstraintAttribute.MESSAGE_TYPE.toString(), "STATE_TRANSITION")
         .addConstraintAttribute(ConstraintAttribute.PARTITION.toString(), ".*")
         .addConstraintAttribute(ConstraintAttribute.CONSTRAINT_VALUE.toString(),
             String.valueOf(PARALLELISM));
     _admin.setConstraint(_clusterName, ConstraintType.MESSAGE_CONSTRAINT, "constraint_1",
         constraintItemBuilder.build());
     runTest();
   }

   /**
    * This test directly embeds the parallelism per partition directly into the upper bound for the
    * ONLINE state. This does not require the controller to support partition-level constraints.
    * @throws Exception
    */
   @Test
   public void testParallelismInStateModel() throws Exception {
     _stateModel = "OnlineOfflineBounded";

     // Add a state model with the parallelism implicit
     _clusterSetup.addStateModelDef(_clusterName, _stateModel,
         createEnforcedParallelismStateModelDef(_stateModel));
     runTest();
   }

   private void runTest() throws Exception {
     final int NUM_PARTITIONS = 1;
     final int NUM_REPLICAS = 2;
     final int NUM_INSTANCES = 2;
     final String RESOURCE_NAME = "MyResource";

     // Setup instances
     String[] instanceInfoArray = {
         "localhost_1", "localhost_2"
     };
     _clusterSetup.addInstancesToCluster(_clusterName, instanceInfoArray);

     // Add resource
     _clusterSetup.addResourceToCluster(_clusterName, RESOURCE_NAME, NUM_PARTITIONS, _stateModel,
         RebalanceMode.SEMI_AUTO.toString());

     // Update resource with empty preference lists
     IdealState idealState = _admin.getResourceIdealState(_clusterName, RESOURCE_NAME);
     for (int i = 0; i < NUM_PARTITIONS; i++) {
       String partitionName = RESOURCE_NAME + "_" + i;
       List<String> dummyPreferences = Lists.newArrayList();
       for (int j = 0; j < NUM_REPLICAS; j++) {
         // hack: need to have some dummy values in the preference list to pass validation
         dummyPreferences.add("");
       }
       idealState.getRecord().setListField(partitionName, dummyPreferences);
     }
     idealState.setReplicas(String.valueOf(NUM_REPLICAS));
     _admin.setResourceIdealState(_clusterName, RESOURCE_NAME, idealState);

     // Start some instances
     HelixManager[] participants = new HelixManager[NUM_INSTANCES];
     for (int i = 0; i < NUM_INSTANCES; i++) {
       participants[i] = HelixManagerFactory.getZKHelixManager(_clusterName, instanceInfoArray[i],
           InstanceType.PARTICIPANT, ZK_ADDR);
       participants[i].getStateMachineEngine().registerStateModelFactory(_stateModel,
           new PrefListTaskOnlineOfflineStateModelFactory());
       participants[i].connect();
     }

     // Start the controller
     HelixManager controller =
         HelixManagerFactory.getZKHelixManager(_clusterName, null, InstanceType.CONTROLLER, ZK_ADDR);
     controller.connect();

     // Disable controller immediately
     _admin.enableCluster(_clusterName, false);

     // This resource only has 1 partition
     String partitionName = RESOURCE_NAME + "_" + 0;

     // There should be no preference lists yet
     Assert.assertTrue(preferenceListIsCorrect(_admin, _clusterName, RESOURCE_NAME, partitionName,
         Arrays.asList("", "")));

     // Add one instance
     addInstanceToPreferences(participants[0].getHelixDataAccessor(),
         participants[0].getInstanceName(), RESOURCE_NAME, Collections.singletonList(partitionName));
     Assert.assertTrue(preferenceListIsCorrect(_admin, _clusterName, RESOURCE_NAME, partitionName,
         Arrays.asList("localhost_1", "")));

     // Add a second instance immediately; the first one should still exist
     addInstanceToPreferences(participants[1].getHelixDataAccessor(),
         participants[1].getInstanceName(), RESOURCE_NAME, Collections.singletonList(partitionName));
     Assert.assertTrue(preferenceListIsCorrect(_admin, _clusterName, RESOURCE_NAME, partitionName,
         Arrays.asList("localhost_1", "localhost_2")));

     // Add the first instance again; it should already exist
     addInstanceToPreferences(participants[0].getHelixDataAccessor(),
         participants[0].getInstanceName(), RESOURCE_NAME, Collections.singletonList(partitionName));
     Assert.assertTrue(preferenceListIsCorrect(_admin, _clusterName, RESOURCE_NAME, partitionName,
         Arrays.asList("localhost_1", "localhost_2")));

     // Prepare for synchronization
     _onlineLatch = new CountDownLatch(2);
     _offlineLatch = new CountDownLatch(2);
     _prefListHistory.clear();

     // Now reenable the controller
     _admin.enableCluster(_clusterName, true);

     // Now wait for both instances to be done
     boolean countReached = _onlineLatch.await(10000, TimeUnit.MILLISECONDS);
     Assert.assertTrue(countReached);
     List<String> top = _prefListHistory.poll();
     assert top != null;
     Assert.assertTrue(top.equals(Arrays.asList("localhost_1", ""))
         || top.equals(Arrays.asList("localhost_2", "")));
     Assert.assertEquals(_prefListHistory.poll(), Arrays.asList("", ""));

     // Wait for everything to be fully offline
     countReached = _offlineLatch.await(10000, TimeUnit.MILLISECONDS);
     Assert.assertTrue(countReached);

     // Add back the instances in the opposite order
     _admin.enableCluster(_clusterName, false);
     addInstanceToPreferences(participants[0].getHelixDataAccessor(),
         participants[1].getInstanceName(), RESOURCE_NAME, Collections.singletonList(partitionName));
     addInstanceToPreferences(participants[0].getHelixDataAccessor(),
         participants[0].getInstanceName(), RESOURCE_NAME, Collections.singletonList(partitionName));
     Assert.assertTrue(preferenceListIsCorrect(_admin, _clusterName, RESOURCE_NAME, partitionName,
         Arrays.asList("localhost_2", "localhost_1")));

     // Reset the latch
     _onlineLatch = new CountDownLatch(2);
     _prefListHistory.clear();
     _admin.enableCluster(_clusterName, true);

     // Now wait both to be done again
     countReached = _onlineLatch.await(10000, TimeUnit.MILLISECONDS);
     Assert.assertTrue(countReached);
     top = _prefListHistory.poll();
     assert top != null;
     Assert.assertTrue(top.equals(Arrays.asList("localhost_1", ""))
         || top.equals(Arrays.asList("localhost_2", "")));
     Assert.assertEquals(_prefListHistory.poll(), Arrays.asList("", ""));
     Assert.assertEquals(_instanceList.size(), 0);

     // Cleanup
     controller.disconnect();
     for (HelixManager participant : participants) {
       participant.disconnect();
     }
   }

   /**
    * Create a modified version of OnlineOffline where the transition to ONLINE is given lowest
    * priority
    * @param stateModelName
    * @return
    */
   private StateModelDefinition createReprioritizedStateModelDef(String stateModelName) {
     StateModelDefinition.Builder builder = new StateModelDefinition.Builder(stateModelName)
         .addState("ONLINE", 1).addState("OFFLINE").addState("DROPPED").addState("ERROR")
         .initialState("OFFLINE").addTransition("ERROR", "OFFLINE", 1)
         .addTransition("ONLINE", "OFFLINE", 2).addTransition("OFFLINE", "DROPPED", 3)
         .addTransition("OFFLINE", "ONLINE", 4).dynamicUpperBound("ONLINE", "R")
         .upperBound("OFFLINE", -1).upperBound("DROPPED", -1).upperBound("ERROR", -1);
     return builder.build();
   }

   /**
    * Create a modified version of OnlineOffline where the parallelism is enforced by the upper bound
    * of ONLINE
    * @param stateModelName
    * @return
    */
   private StateModelDefinition createEnforcedParallelismStateModelDef(String stateModelName) {
     StateModelDefinition.Builder builder =
         new StateModelDefinition.Builder(stateModelName).addState("ONLINE", 1).addState("OFFLINE")
             .addState("DROPPED").addState("ERROR").initialState("OFFLINE")
             .addTransition("ERROR", "OFFLINE", 1).addTransition("ONLINE", "OFFLINE", 2)
             .addTransition("OFFLINE", "DROPPED", 3).addTransition("OFFLINE", "ONLINE", 4)
             .dynamicUpperBound("ONLINE", String.valueOf(PARALLELISM)).upperBound("OFFLINE", -1)
             .upperBound("DROPPED", -1).upperBound("ERROR", -1);
     return builder.build();
   }

   /**
    * Check if the provided list matches the currently persisted preference list
    * @param admin
    * @param clusterName
    * @param resourceName
    * @param partitionName
    * @param expectPreferenceList
    * @return
    */
   private boolean preferenceListIsCorrect(HelixAdmin admin, String clusterName, String resourceName,
       String partitionName, List<String> expectPreferenceList) {
     IdealState idealState = admin.getResourceIdealState(clusterName, resourceName);
     List<String> preferenceList = idealState.getPreferenceList(partitionName);
     return expectPreferenceList.equals(preferenceList);
   }

   /**
    * Update an ideal state so that partitions will have an instance removed from their preference
    * lists
    * @param accessor
    * @param instanceName
    * @param resourceName
    * @param partitionName
    */
   private void removeInstanceFromPreferences(HelixDataAccessor accessor, final String instanceName,
       final String resourceName, final String partitionName) {
     PropertyKey.Builder keyBuilder = accessor.keyBuilder();
     String idealStatePath = keyBuilder.idealStates(resourceName).getPath();
     synchronized (_prefListHistory) {
       // Updater for ideal state
       final List<String> prefList = Lists.newLinkedList();
       DataUpdater<ZNRecord> idealStateUpdater = currentData -> {
         List<String> preferenceList = currentData.getListField(partitionName);
         int numReplicas =
             Integer.valueOf(currentData.getSimpleField(IdealStateProperty.REPLICAS.toString()));
         List<String> newPrefList =
             removeInstanceFromPreferenceList(preferenceList, instanceName, numReplicas);
         currentData.setListField(partitionName, newPrefList);
         prefList.clear();
         prefList.addAll(newPrefList);
         return currentData;
       };
       List<DataUpdater<ZNRecord>> updaters = Lists.newArrayList();
       updaters.add(idealStateUpdater);
       accessor.updateChildren(Collections.singletonList(idealStatePath), updaters,
           AccessOption.PERSISTENT);
       _prefListHistory.add(prefList);
     }
   }

   /**
    * Update an ideal state so that partitions will have a new instance at the tails of their
    * preference lists
    * @param accessor
    * @param instanceName
    * @param resourceName
    * @param partitions
    */
   private void addInstanceToPreferences(HelixDataAccessor accessor, final String instanceName,
       final String resourceName, final List<String> partitions) {
     PropertyKey.Builder keyBuilder = accessor.keyBuilder();
     String idealStatePath = keyBuilder.idealStates(resourceName).getPath();
     synchronized (_prefListHistory) {
       // Updater for ideal state
       final List<String> prefList = Lists.newLinkedList();
       DataUpdater<ZNRecord> idealStateUpdater = currentData -> {
         for (String partitionName : partitions) {
           List<String> preferenceList = currentData.getListField(partitionName);
           int numReplicas =
               Integer.valueOf(currentData.getSimpleField(IdealStateProperty.REPLICAS.toString()));
           List<String> newPrefList =
               addInstanceToPreferenceList(preferenceList, instanceName, numReplicas);
           currentData.setListField(partitionName, newPrefList);
           prefList.clear();
           prefList.addAll(newPrefList);
         }
         return currentData;
       };

       // Send update requests together
       List<DataUpdater<ZNRecord>> updaters = Lists.newArrayList();
       updaters.add(idealStateUpdater);
       accessor.updateChildren(Collections.singletonList(idealStatePath), updaters,
           AccessOption.PERSISTENT);
       _prefListHistory.add(prefList);
     }
   }

   /**
    * Add the instance to the preference list, removing padding as necessary
    * @param existPreferences
    * @param instanceName
    * @param numReplicas
    * @return A new preference list with the instance added to the end (if it did not already exist)
    */
   private static List<String> addInstanceToPreferenceList(List<String> existPreferences,
       String instanceName, int numReplicas) {
     if (existPreferences == null) {
       existPreferences = Lists.newArrayList();
     }
     List<String> newPreferences = Lists.newArrayListWithCapacity(numReplicas);
     for (String existInstance : existPreferences) {
       if (!existInstance.isEmpty()) {
         newPreferences.add(existInstance);
       }
     }
     if (!newPreferences.contains(instanceName)) {
       newPreferences.add(instanceName);
     }
     addDummiesToPreferenceList(newPreferences, numReplicas);
     return newPreferences;
   }

   /**
    * Remove an instance from a preference list, padding as necessary
    * @param existPreferences
    * @param instanceName
    * @param numReplicas
    * @return new preference list with the instance removed (if it existed)
    */
   private static List<String> removeInstanceFromPreferenceList(List<String> existPreferences,
       String instanceName, int numReplicas) {
     if (existPreferences == null) {
       existPreferences = Lists.newArrayList();
     }
     List<String> newPreferences = Lists.newArrayListWithCapacity(numReplicas);
     for (String existInstance : existPreferences) {
       if (!existInstance.isEmpty() && !existInstance.equals(instanceName)) {
         newPreferences.add(existInstance);
       }
     }
     addDummiesToPreferenceList(newPreferences, numReplicas);
     return newPreferences;
   }

   /**
    * Pad a preference list with empty strings to have it reach numReplicas size
    * @param preferences
    * @param numReplicas
    */
   private static void addDummiesToPreferenceList(List<String> preferences, int numReplicas) {
     int numRemaining = numReplicas - preferences.size();
     numRemaining = (numRemaining > 0) ? numRemaining : 0;
     for (int i = 0; i < numRemaining; i++) {
       preferences.add("");
     }
   }

   /**
    * A state model whose callbacks correspond to when to run a task. When becoming online, the task
    * should be run, and then the instance should be removed from the preference list for the task
    * (padded by spaces). All other transitions are no-ops.
    */
   public class PrefListTaskOnlineOfflineStateModel extends StateModel {
     /**
      * Run the task. The parallelism of this is dictated by the constraints that are set.
      * @param message
      * @param context
      * @throws InterruptedException
      */
     public void onBecomeOnlineFromOffline(final Message message, NotificationContext context)
         throws InterruptedException {
       // Do the work, and then finally remove the instance from the preference list for this
       // partition
       HelixManager manager = context.getManager();
       LOG.info("START onBecomeOnlineFromOffline for " + message.getPartitionName() + " on "
           + manager.getInstanceName());
       int oldSize;
       synchronized (_instanceList) {
         oldSize = _instanceList.size();
         _instanceList.add(manager.getInstanceName());
       }
       Assert.assertEquals(oldSize, 0); // ensure these transitions are fully synchronized

       Thread.sleep(TRANSITION_TIME); // a sleep simulates work

       // Need to disable in order to get the transition the next time
       HelixDataAccessor accessor = manager.getHelixDataAccessor();
       removeInstanceFromPreferences(accessor, manager.getInstanceName(), message.getResourceName(),
           message.getPartitionName());
       LOG.info("FINISH onBecomeOnlineFromOffline for " + message.getPartitionName() + " on "
           + manager.getInstanceName());

       int newSize;
       synchronized (_instanceList) {
         _instanceList.remove(_instanceList.size() - 1);
         newSize = _instanceList.size();
       }
       Assert.assertEquals(newSize, oldSize); // ensure nothing came in during this time
       _onlineLatch.countDown();
     }

     public void onBecomeOfflineFromOnline(Message message, NotificationContext context) {
       HelixManager manager = context.getManager();
       LOG.info("onBecomeOfflineFromOnline for " + message.getPartitionName() + " on "
           + manager.getInstanceName());
     }

     public void onBecomeDroppedFromOffline(Message message, NotificationContext context) {
       HelixManager manager = context.getManager();
       LOG.info("onBecomeDroppedFromOffline for " + message.getPartitionName() + " on "
           + manager.getInstanceName());
       _offlineLatch.countDown();
     }

     public void onBecomeOfflineFromError(Message message, NotificationContext context) {
       HelixManager manager = context.getManager();
       LOG.info("onBecomeOfflineFromError for " + message.getPartitionName() + " on "
           + manager.getInstanceName());
     }
   }

   public class PrefListTaskOnlineOfflineStateModelFactory
       extends StateModelFactory<PrefListTaskOnlineOfflineStateModel> {
     @Override
     public PrefListTaskOnlineOfflineStateModel createNewStateModel(String resource,
         String partitionName) {
       return new PrefListTaskOnlineOfflineStateModel();
     }
   }
 }
	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.Arrays;
	import java.util.Collections;
	import java.util.List;
	import java.util.Queue;
	import java.util.concurrent.CountDownLatch;
	import java.util.concurrent.TimeUnit;

	import com.google.common.collect.Lists;
	import org.apache.helix.AccessOption;
	import org.apache.helix.HelixAdmin;
	import org.apache.helix.HelixDataAccessor;
	import org.apache.helix.HelixManager;
	import org.apache.helix.HelixManagerFactory;
	import org.apache.helix.InstanceType;
	import org.apache.helix.NotificationContext;
	import org.apache.helix.PropertyKey;
	import org.apache.helix.TestHelper;
	import org.apache.helix.zookeeper.datamodel.ZNRecord;
	import org.apache.helix.ZkUnitTestBase;
	import org.apache.helix.model.ClusterConstraints.ConstraintAttribute;
	import org.apache.helix.model.ClusterConstraints.ConstraintType;
	import org.apache.helix.model.IdealState;
	import org.apache.helix.model.IdealState.IdealStateProperty;
	import org.apache.helix.model.IdealState.RebalanceMode;
	import org.apache.helix.model.Message;
	import org.apache.helix.model.StateModelDefinition;
	import org.apache.helix.model.builder.ConstraintItemBuilder;
	import org.apache.helix.participant.statemachine.StateModel;
	import org.apache.helix.participant.statemachine.StateModelFactory;
	import org.apache.helix.tools.ClusterSetup;
	import org.apache.helix.zookeeper.zkclient.DataUpdater;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;
	import org.testng.Assert;
	import org.testng.annotations.AfterClass;
	import org.testng.annotations.BeforeMethod;
	import org.testng.annotations.Test;

	public class TestPreferenceListAsQueue extends ZkUnitTestBase {
	private static final Logger LOG = LoggerFactory.getLogger(TestPreferenceListAsQueue.class);
	private static final int TRANSITION_TIME = 500;
	private static final int PARALLELISM = 1;

	private List<String> _instanceList;
	private Queue<List<String>> _prefListHistory;
	private String _clusterName;
	private String _stateModel;
	private ClusterSetup _clusterSetup;
	private HelixAdmin _admin;
	private CountDownLatch _onlineLatch;
	private CountDownLatch _offlineLatch;

	@BeforeMethod
	public void beforeMethod() {
	_instanceList = Lists.newLinkedList();
	_clusterSetup = new ClusterSetup(ZK_ADDR);
	_admin = _clusterSetup.getClusterManagementTool();
	_prefListHistory = Lists.newLinkedList();

	// Create cluster
	String className = TestHelper.getTestClassName();
	String methodName = TestHelper.getTestMethodName();
	_clusterName = className + "_" + methodName;
	_clusterSetup.addCluster(_clusterName, true);
	}

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

	/**
	* This test ensures successful creation when the state model has OFFLINE --> deprioritized and
	* a partition-level constraint enforces parallelism
	* @throws Exception
	*/
	@Test
	public void testReprioritizedWithConstraint() throws Exception {
	_stateModel = "OnlineOfflineReprioritized";

	// Add a state model with the transition to ONLINE deprioritized
	_clusterSetup.addStateModelDef(_clusterName, _stateModel,
	createReprioritizedStateModelDef(_stateModel));

	// Add a constraint of one transition per partition
	ConstraintItemBuilder constraintItemBuilder = new ConstraintItemBuilder();
	constraintItemBuilder
	.addConstraintAttribute(ConstraintAttribute.MESSAGE_TYPE.toString(), "STATE_TRANSITION")
	.addConstraintAttribute(ConstraintAttribute.PARTITION.toString(), ".*")
	.addConstraintAttribute(ConstraintAttribute.CONSTRAINT_VALUE.toString(),
	String.valueOf(PARALLELISM));
	_admin.setConstraint(_clusterName, ConstraintType.MESSAGE_CONSTRAINT, "constraint_1",
	constraintItemBuilder.build());
	runTest();
	}

	/**
	* This test directly embeds the parallelism per partition directly into the upper bound for the
	* ONLINE state. This does not require the controller to support partition-level constraints.
	* @throws Exception
	*/
	@Test
	public void testParallelismInStateModel() throws Exception {
	_stateModel = "OnlineOfflineBounded";

	// Add a state model with the parallelism implicit
	_clusterSetup.addStateModelDef(_clusterName, _stateModel,
	createEnforcedParallelismStateModelDef(_stateModel));
	runTest();
	}

	private void runTest() throws Exception {
	final int NUM_PARTITIONS = 1;
	final int NUM_REPLICAS = 2;
	final int NUM_INSTANCES = 2;
	final String RESOURCE_NAME = "MyResource";

	// Setup instances
	String[] instanceInfoArray = {
	"localhost_1", "localhost_2"
	};
	_clusterSetup.addInstancesToCluster(_clusterName, instanceInfoArray);

	// Add resource
	_clusterSetup.addResourceToCluster(_clusterName, RESOURCE_NAME, NUM_PARTITIONS, _stateModel,
	RebalanceMode.SEMI_AUTO.toString());

	// Update resource with empty preference lists
	IdealState idealState = _admin.getResourceIdealState(_clusterName, RESOURCE_NAME);
	for (int i = 0; i < NUM_PARTITIONS; i++) {
	String partitionName = RESOURCE_NAME + "_" + i;
	List<String> dummyPreferences = Lists.newArrayList();
	for (int j = 0; j < NUM_REPLICAS; j++) {
	// hack: need to have some dummy values in the preference list to pass validation
	dummyPreferences.add("");
	}
	idealState.getRecord().setListField(partitionName, dummyPreferences);
	}
	idealState.setReplicas(String.valueOf(NUM_REPLICAS));
	_admin.setResourceIdealState(_clusterName, RESOURCE_NAME, idealState);

	// Start some instances
	HelixManager[] participants = new HelixManager[NUM_INSTANCES];
	for (int i = 0; i < NUM_INSTANCES; i++) {
	participants[i] = HelixManagerFactory.getZKHelixManager(_clusterName, instanceInfoArray[i],
	InstanceType.PARTICIPANT, ZK_ADDR);
	participants[i].getStateMachineEngine().registerStateModelFactory(_stateModel,
	new PrefListTaskOnlineOfflineStateModelFactory());
	participants[i].connect();
	}

	// Start the controller
	HelixManager controller =
	HelixManagerFactory.getZKHelixManager(_clusterName, null, InstanceType.CONTROLLER, ZK_ADDR);
	controller.connect();

	// Disable controller immediately
	_admin.enableCluster(_clusterName, false);

	// This resource only has 1 partition
	String partitionName = RESOURCE_NAME + "_" + 0;

	// There should be no preference lists yet
	Assert.assertTrue(preferenceListIsCorrect(_admin, _clusterName, RESOURCE_NAME, partitionName,
	Arrays.asList("", "")));

	// Add one instance
	addInstanceToPreferences(participants[0].getHelixDataAccessor(),
	participants[0].getInstanceName(), RESOURCE_NAME, Collections.singletonList(partitionName));
	Assert.assertTrue(preferenceListIsCorrect(_admin, _clusterName, RESOURCE_NAME, partitionName,
	Arrays.asList("localhost_1", "")));

	// Add a second instance immediately; the first one should still exist
	addInstanceToPreferences(participants[1].getHelixDataAccessor(),
	participants[1].getInstanceName(), RESOURCE_NAME, Collections.singletonList(partitionName));
	Assert.assertTrue(preferenceListIsCorrect(_admin, _clusterName, RESOURCE_NAME, partitionName,
	Arrays.asList("localhost_1", "localhost_2")));

	// Add the first instance again; it should already exist
	addInstanceToPreferences(participants[0].getHelixDataAccessor(),
	participants[0].getInstanceName(), RESOURCE_NAME, Collections.singletonList(partitionName));
	Assert.assertTrue(preferenceListIsCorrect(_admin, _clusterName, RESOURCE_NAME, partitionName,
	Arrays.asList("localhost_1", "localhost_2")));

	// Prepare for synchronization
	_onlineLatch = new CountDownLatch(2);
	_offlineLatch = new CountDownLatch(2);
	_prefListHistory.clear();

	// Now reenable the controller
	_admin.enableCluster(_clusterName, true);

	// Now wait for both instances to be done
	boolean countReached = _onlineLatch.await(10000, TimeUnit.MILLISECONDS);
	Assert.assertTrue(countReached);
	List<String> top = _prefListHistory.poll();
	assert top != null;
	Assert.assertTrue(top.equals(Arrays.asList("localhost_1", ""))
	\|\| top.equals(Arrays.asList("localhost_2", "")));
	Assert.assertEquals(_prefListHistory.poll(), Arrays.asList("", ""));

	// Wait for everything to be fully offline
	countReached = _offlineLatch.await(10000, TimeUnit.MILLISECONDS);
	Assert.assertTrue(countReached);

	// Add back the instances in the opposite order
	_admin.enableCluster(_clusterName, false);
	addInstanceToPreferences(participants[0].getHelixDataAccessor(),
	participants[1].getInstanceName(), RESOURCE_NAME, Collections.singletonList(partitionName));
	addInstanceToPreferences(participants[0].getHelixDataAccessor(),
	participants[0].getInstanceName(), RESOURCE_NAME, Collections.singletonList(partitionName));
	Assert.assertTrue(preferenceListIsCorrect(_admin, _clusterName, RESOURCE_NAME, partitionName,
	Arrays.asList("localhost_2", "localhost_1")));

	// Reset the latch
	_onlineLatch = new CountDownLatch(2);
	_prefListHistory.clear();
	_admin.enableCluster(_clusterName, true);

	// Now wait both to be done again
	countReached = _onlineLatch.await(10000, TimeUnit.MILLISECONDS);
	Assert.assertTrue(countReached);
	top = _prefListHistory.poll();
	assert top != null;
	Assert.assertTrue(top.equals(Arrays.asList("localhost_1", ""))
	\|\| top.equals(Arrays.asList("localhost_2", "")));
	Assert.assertEquals(_prefListHistory.poll(), Arrays.asList("", ""));
	Assert.assertEquals(_instanceList.size(), 0);

	// Cleanup
	controller.disconnect();
	for (HelixManager participant : participants) {
	participant.disconnect();
	}
	}

	/**
	* Create a modified version of OnlineOffline where the transition to ONLINE is given lowest
	* priority
	* @param stateModelName
	* @return
	*/
	private StateModelDefinition createReprioritizedStateModelDef(String stateModelName) {
	StateModelDefinition.Builder builder = new StateModelDefinition.Builder(stateModelName)
	.addState("ONLINE", 1).addState("OFFLINE").addState("DROPPED").addState("ERROR")
	.initialState("OFFLINE").addTransition("ERROR", "OFFLINE", 1)
	.addTransition("ONLINE", "OFFLINE", 2).addTransition("OFFLINE", "DROPPED", 3)
	.addTransition("OFFLINE", "ONLINE", 4).dynamicUpperBound("ONLINE", "R")
	.upperBound("OFFLINE", -1).upperBound("DROPPED", -1).upperBound("ERROR", -1);
	return builder.build();
	}

	/**
	* Create a modified version of OnlineOffline where the parallelism is enforced by the upper bound
	* of ONLINE
	* @param stateModelName
	* @return
	*/
	private StateModelDefinition createEnforcedParallelismStateModelDef(String stateModelName) {
	StateModelDefinition.Builder builder =
	new StateModelDefinition.Builder(stateModelName).addState("ONLINE", 1).addState("OFFLINE")
	.addState("DROPPED").addState("ERROR").initialState("OFFLINE")
	.addTransition("ERROR", "OFFLINE", 1).addTransition("ONLINE", "OFFLINE", 2)
	.addTransition("OFFLINE", "DROPPED", 3).addTransition("OFFLINE", "ONLINE", 4)
	.dynamicUpperBound("ONLINE", String.valueOf(PARALLELISM)).upperBound("OFFLINE", -1)
	.upperBound("DROPPED", -1).upperBound("ERROR", -1);
	return builder.build();
	}

	/**
	* Check if the provided list matches the currently persisted preference list
	* @param admin
	* @param clusterName
	* @param resourceName
	* @param partitionName
	* @param expectPreferenceList
	* @return
	*/
	private boolean preferenceListIsCorrect(HelixAdmin admin, String clusterName, String resourceName,
	String partitionName, List<String> expectPreferenceList) {
	IdealState idealState = admin.getResourceIdealState(clusterName, resourceName);
	List<String> preferenceList = idealState.getPreferenceList(partitionName);
	return expectPreferenceList.equals(preferenceList);
	}

	/**
	* Update an ideal state so that partitions will have an instance removed from their preference
	* lists
	* @param accessor
	* @param instanceName
	* @param resourceName
	* @param partitionName
	*/
	private void removeInstanceFromPreferences(HelixDataAccessor accessor, final String instanceName,
	final String resourceName, final String partitionName) {
	PropertyKey.Builder keyBuilder = accessor.keyBuilder();
	String idealStatePath = keyBuilder.idealStates(resourceName).getPath();
	synchronized (_prefListHistory) {
	// Updater for ideal state
	final List<String> prefList = Lists.newLinkedList();
	DataUpdater<ZNRecord> idealStateUpdater = currentData -> {
	List<String> preferenceList = currentData.getListField(partitionName);
	int numReplicas =
	Integer.valueOf(currentData.getSimpleField(IdealStateProperty.REPLICAS.toString()));
	List<String> newPrefList =
	removeInstanceFromPreferenceList(preferenceList, instanceName, numReplicas);
	currentData.setListField(partitionName, newPrefList);
	prefList.clear();
	prefList.addAll(newPrefList);
	return currentData;
	};
	List<DataUpdater<ZNRecord>> updaters = Lists.newArrayList();
	updaters.add(idealStateUpdater);
	accessor.updateChildren(Collections.singletonList(idealStatePath), updaters,
	AccessOption.PERSISTENT);
	_prefListHistory.add(prefList);
	}
	}

	/**
	* Update an ideal state so that partitions will have a new instance at the tails of their
	* preference lists
	* @param accessor
	* @param instanceName
	* @param resourceName
	* @param partitions
	*/
	private void addInstanceToPreferences(HelixDataAccessor accessor, final String instanceName,
	final String resourceName, final List<String> partitions) {
	PropertyKey.Builder keyBuilder = accessor.keyBuilder();
	String idealStatePath = keyBuilder.idealStates(resourceName).getPath();
	synchronized (_prefListHistory) {
	// Updater for ideal state
	final List<String> prefList = Lists.newLinkedList();
	DataUpdater<ZNRecord> idealStateUpdater = currentData -> {
	for (String partitionName : partitions) {
	List<String> preferenceList = currentData.getListField(partitionName);
	int numReplicas =
	Integer.valueOf(currentData.getSimpleField(IdealStateProperty.REPLICAS.toString()));
	List<String> newPrefList =
	addInstanceToPreferenceList(preferenceList, instanceName, numReplicas);
	currentData.setListField(partitionName, newPrefList);
	prefList.clear();
	prefList.addAll(newPrefList);
	}
	return currentData;
	};

	// Send update requests together
	List<DataUpdater<ZNRecord>> updaters = Lists.newArrayList();
	updaters.add(idealStateUpdater);
	accessor.updateChildren(Collections.singletonList(idealStatePath), updaters,
	AccessOption.PERSISTENT);
	_prefListHistory.add(prefList);
	}
	}

	/**
	* Add the instance to the preference list, removing padding as necessary
	* @param existPreferences
	* @param instanceName
	* @param numReplicas
	* @return A new preference list with the instance added to the end (if it did not already exist)
	*/
	private static List<String> addInstanceToPreferenceList(List<String> existPreferences,
	String instanceName, int numReplicas) {
	if (existPreferences == null) {
	existPreferences = Lists.newArrayList();
	}
	List<String> newPreferences = Lists.newArrayListWithCapacity(numReplicas);
	for (String existInstance : existPreferences) {
	if (!existInstance.isEmpty()) {
	newPreferences.add(existInstance);
	}
	}
	if (!newPreferences.contains(instanceName)) {
	newPreferences.add(instanceName);
	}
	addDummiesToPreferenceList(newPreferences, numReplicas);
	return newPreferences;
	}

	/**
	* Remove an instance from a preference list, padding as necessary
	* @param existPreferences
	* @param instanceName
	* @param numReplicas
	* @return new preference list with the instance removed (if it existed)
	*/
	private static List<String> removeInstanceFromPreferenceList(List<String> existPreferences,
	String instanceName, int numReplicas) {
	if (existPreferences == null) {
	existPreferences = Lists.newArrayList();
	}
	List<String> newPreferences = Lists.newArrayListWithCapacity(numReplicas);
	for (String existInstance : existPreferences) {
	if (!existInstance.isEmpty() && !existInstance.equals(instanceName)) {
	newPreferences.add(existInstance);
	}
	}
	addDummiesToPreferenceList(newPreferences, numReplicas);
	return newPreferences;
	}

	/**
	* Pad a preference list with empty strings to have it reach numReplicas size
	* @param preferences
	* @param numReplicas
	*/
	private static void addDummiesToPreferenceList(List<String> preferences, int numReplicas) {
	int numRemaining = numReplicas - preferences.size();
	numRemaining = (numRemaining > 0) ? numRemaining : 0;
	for (int i = 0; i < numRemaining; i++) {
	preferences.add("");
	}
	}

	/**
	* A state model whose callbacks correspond to when to run a task. When becoming online, the task
	* should be run, and then the instance should be removed from the preference list for the task
	* (padded by spaces). All other transitions are no-ops.
	*/
	public class PrefListTaskOnlineOfflineStateModel extends StateModel {
	/**
	* Run the task. The parallelism of this is dictated by the constraints that are set.
	* @param message
	* @param context
	* @throws InterruptedException
	*/
	public void onBecomeOnlineFromOffline(final Message message, NotificationContext context)
	throws InterruptedException {
	// Do the work, and then finally remove the instance from the preference list for this
	// partition
	HelixManager manager = context.getManager();
	LOG.info("START onBecomeOnlineFromOffline for " + message.getPartitionName() + " on "
	+ manager.getInstanceName());
	int oldSize;
	synchronized (_instanceList) {
	oldSize = _instanceList.size();
	_instanceList.add(manager.getInstanceName());
	}
	Assert.assertEquals(oldSize, 0); // ensure these transitions are fully synchronized

	Thread.sleep(TRANSITION_TIME); // a sleep simulates work

	// Need to disable in order to get the transition the next time
	HelixDataAccessor accessor = manager.getHelixDataAccessor();
	removeInstanceFromPreferences(accessor, manager.getInstanceName(), message.getResourceName(),
	message.getPartitionName());
	LOG.info("FINISH onBecomeOnlineFromOffline for " + message.getPartitionName() + " on "
	+ manager.getInstanceName());

	int newSize;
	synchronized (_instanceList) {
	_instanceList.remove(_instanceList.size() - 1);
	newSize = _instanceList.size();
	}
	Assert.assertEquals(newSize, oldSize); // ensure nothing came in during this time
	_onlineLatch.countDown();
	}

	public void onBecomeOfflineFromOnline(Message message, NotificationContext context) {
	HelixManager manager = context.getManager();
	LOG.info("onBecomeOfflineFromOnline for " + message.getPartitionName() + " on "
	+ manager.getInstanceName());
	}

	public void onBecomeDroppedFromOffline(Message message, NotificationContext context) {
	HelixManager manager = context.getManager();
	LOG.info("onBecomeDroppedFromOffline for " + message.getPartitionName() + " on "
	+ manager.getInstanceName());
	_offlineLatch.countDown();
	}

	public void onBecomeOfflineFromError(Message message, NotificationContext context) {
	HelixManager manager = context.getManager();
	LOG.info("onBecomeOfflineFromError for " + message.getPartitionName() + " on "
	+ manager.getInstanceName());
	}
	}

	public class PrefListTaskOnlineOfflineStateModelFactory
	extends StateModelFactory<PrefListTaskOnlineOfflineStateModel> {
	@Override
	public PrefListTaskOnlineOfflineStateModel createNewStateModel(String resource,
	String partitionName) {
	return new PrefListTaskOnlineOfflineStateModel();
	}
	}
	}