helix-core/src/test/java/org/apache/helix/messaging/p2pMessage/TestP2PMessages.java - helix - Git at Google

 package org.apache.helix.messaging.p2pMessage;

 /*
  * 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.lang.reflect.Method;
 import java.util.Date;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.Executors;

 import org.apache.helix.HelixConstants;
 import org.apache.helix.HelixDefinedState;
 import org.apache.helix.PropertyKey;
 import org.apache.helix.controller.common.ResourcesStateMap;
 import org.apache.helix.controller.dataproviders.ResourceControllerDataProvider;
 import org.apache.helix.controller.pipeline.Pipeline;
 import org.apache.helix.controller.stages.AttributeName;
 import org.apache.helix.controller.stages.BaseStageTest;
 import org.apache.helix.controller.stages.BestPossibleStateCalcStage;
 import org.apache.helix.controller.stages.CurrentStateComputationStage;
 import org.apache.helix.controller.stages.IntermediateStateCalcStage;
 import org.apache.helix.controller.stages.MessageGenerationPhase;
 import org.apache.helix.controller.stages.MessageSelectionStage;
 import org.apache.helix.controller.stages.MessageThrottleStage;
 import org.apache.helix.controller.stages.ReadClusterDataStage;
 import org.apache.helix.controller.stages.ResourceComputationStage;
 import org.apache.helix.controller.stages.resource.ResourceMessageDispatchStage;
 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.MasterSlaveSMD;
 import org.apache.helix.model.Message;
 import org.apache.helix.model.Partition;
 import org.testng.Assert;
 import org.testng.ITestContext;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;

 public class TestP2PMessages extends BaseStageTest {
   private String _db = "testDB";
   private int _numPartition = 1;
   private int _numReplica = 3;

   private Partition _partition = new Partition(_db + "_0");

   private ResourceControllerDataProvider _dataCache;
   private Pipeline _fullPipeline;

   private ResourcesStateMap _initialStateMap;

   private Set<String> _instances;
   private Map<String, LiveInstance> _liveInstanceMap;
   private String _initialMaster;

   @BeforeClass
   public void beforeClass() {
     super.beforeClass();
     setup();
     setupIdealState(3, new String[] { _db }, _numPartition, _numReplica,
         IdealState.RebalanceMode.SEMI_AUTO, BuiltInStateModelDefinitions.MasterSlave.name());
     setupStateModel();
     setupInstances(3);
     setupLiveInstances(3);

     ClusterConfig clusterConfig = new ClusterConfig(_clusterName);
     clusterConfig.enableP2PMessage(true);
     setClusterConfig(clusterConfig);

     _dataCache = new ResourceControllerDataProvider(_clusterName);
     _dataCache.setAsyncTasksThreadPool(Executors.newSingleThreadExecutor());

     _dataCache.refresh(manager.getHelixDataAccessor());

     event.addAttribute(AttributeName.ControllerDataProvider.name(), _dataCache);
     event.addAttribute(AttributeName.helixmanager.name(), manager);

     _fullPipeline = new Pipeline("FullPipeline");
     _fullPipeline.addStage(new ReadClusterDataStage());
     _fullPipeline.addStage(new ResourceComputationStage());
     _fullPipeline.addStage(new CurrentStateComputationStage());
     _fullPipeline.addStage(new BestPossibleStateCalcStage());
     _fullPipeline.addStage(new MessageGenerationPhase());
     _fullPipeline.addStage(new MessageSelectionStage());
     _fullPipeline.addStage(new IntermediateStateCalcStage());
     _fullPipeline.addStage(new MessageThrottleStage());
     _fullPipeline.addStage(new ResourceMessageDispatchStage());

     try {
       _fullPipeline.handle(event);
     } catch (Exception e) {
       e.printStackTrace();
     }

     _instances = _dataCache.getAllInstances();
     _liveInstanceMap = _dataCache.getLiveInstances();

     _initialStateMap = event.getAttribute(AttributeName.BEST_POSSIBLE_STATE.name());
     _initialMaster = getTopStateInstance(_initialStateMap.getInstanceStateMap(_db, _partition),
         MasterSlaveSMD.States.MASTER.name());
     Assert.assertNotNull(_initialMaster);
   }

   @BeforeMethod  // just to overide the per-test setup in base class.
   public void beforeTest(Method testMethod, ITestContext testContext) {
     long startTime = System.currentTimeMillis();
     System.out.println("START " + testMethod.getName() + " at " + new Date(startTime));
     testContext.setAttribute("StartTime", System.currentTimeMillis());
   }

   @Test
   public void testP2PSendAndTimeout() throws Exception {
     reset(_initialStateMap);

     // Disable old master ((initialMaster) instance,
     // Validate: a M->S message should be sent to initialMaster with a P2P message attached for secondMaster.
     admin.enableInstance(_clusterName, _initialMaster, false);
     _dataCache = event.getAttribute(AttributeName.ControllerDataProvider.name());
     _dataCache.notifyDataChange(HelixConstants.ChangeType.INSTANCE_CONFIG);

     _fullPipeline.handle(event);

     ResourcesStateMap bestpossibleState =
         event.getAttribute(AttributeName.INTERMEDIATE_STATE.name());

     List<Message> messages = getMessages(_initialMaster);

     Assert.assertEquals(messages.size(), 1);
     Message toSlaveMessage = messages.get(0);
     Assert.assertEquals(toSlaveMessage.getTgtName(), _initialMaster);
     Assert.assertEquals(toSlaveMessage.getFromState(), MasterSlaveSMD.States.MASTER.name());
     Assert.assertEquals(toSlaveMessage.getToState(), MasterSlaveSMD.States.SLAVE.name());

     // verify p2p message are attached to the M->S message sent to the old master instance
     Assert.assertEquals(toSlaveMessage.getRelayMessages().entrySet().size(), 1);
     String secondMaster =
         getTopStateInstance(bestpossibleState.getInstanceStateMap(_db, _partition),
             MasterSlaveSMD.States.MASTER.name());

     Message relayMessage = toSlaveMessage.getRelayMessage(secondMaster);
     Assert.assertNotNull(relayMessage);
     Assert.assertEquals(relayMessage.getMsgSubType(), Message.MessageType.RELAYED_MESSAGE.name());
     Assert.assertEquals(relayMessage.getTgtName(), secondMaster);
     Assert.assertEquals(relayMessage.getRelaySrcHost(), _initialMaster);
     Assert.assertEquals(relayMessage.getFromState(), MasterSlaveSMD.States.SLAVE.name());
     Assert.assertEquals(relayMessage.getToState(), MasterSlaveSMD.States.MASTER.name());

     // Old master (initialMaster) completed the M->S transition,
     // but has not forward p2p message to new master (secondMaster) yet.
     // Validate: Controller should not send S->M message to new master.
     handleMessage(_initialMaster, _db);
     _fullPipeline.handle(event);
     messages = getMessages(secondMaster);
     Assert.assertEquals(messages.size(), 0);

     // Old master (initialMaster) completed the M->S transition,
     // but has not forward p2p message to new master (secondMaster) yet, but p2p message should already timeout.
     // Validate: Controller should send S->M message to new master.
     Thread.sleep(Message.RELAY_MESSAGE_DEFAULT_EXPIRY);

     _fullPipeline.handle(event);
     messages = getMessages(secondMaster);
     Assert.assertEquals(messages.size(), 1);
     Assert.assertEquals(messages.get(0).getTgtName(), secondMaster);
     Assert.assertEquals(messages.get(0).getFromState(), MasterSlaveSMD.States.SLAVE.name());
     Assert.assertEquals(messages.get(0).getToState(), MasterSlaveSMD.States.MASTER.name());
   }

   @Test
   public void testP2PWithErrorState() throws Exception {
     reset(_initialStateMap);
     // Disable old master ((initialMaster) instance,
     // Validate: a M->S message should be sent to initialMaster with a P2P message attached for secondMaster.

     // disable existing master instance
     admin.enableInstance(_clusterName, _initialMaster, false);
     _dataCache.notifyDataChange(HelixConstants.ChangeType.INSTANCE_CONFIG);
     _fullPipeline.handle(event);

     ResourcesStateMap bestpossibleState =
         event.getAttribute(AttributeName.INTERMEDIATE_STATE.name());

     List<Message> messages = getMessages(_initialMaster);

     Assert.assertEquals(messages.size(), 1);
     Message toSlaveMessage = messages.get(0);

     // verify p2p message are attached to the M->S message sent to the old master instance
     Assert.assertEquals(toSlaveMessage.getRelayMessages().entrySet().size(), 1);
     String secondMaster =
         getTopStateInstance(bestpossibleState.getInstanceStateMap(_db, _partition),
             MasterSlaveSMD.States.MASTER.name());
     Message relayMessage = toSlaveMessage.getRelayMessage(secondMaster);
     Assert.assertNotNull(relayMessage);
     Assert.assertEquals(relayMessage.getMsgSubType(), Message.MessageType.RELAYED_MESSAGE.name());
     Assert.assertEquals(relayMessage.getTgtName(), secondMaster);

     // Old master (initialMaster) failed the M->S transition,
     // but has not forward p2p message to new master (secondMaster) yet.
     // Validate: Controller should ignore the ERROR partition and send S->M message to new master.
     String session = _dataCache.getLiveInstances().get(_initialMaster).getEphemeralOwner();
     PropertyKey currentStateKey =
         new PropertyKey.Builder(_clusterName).currentState(_initialMaster, session, _db);
     CurrentState currentState = accessor.getProperty(currentStateKey);
     currentState
         .setPreviousState(_partition.getPartitionName(), MasterSlaveSMD.States.MASTER.name());
     currentState.setState(_partition.getPartitionName(), HelixDefinedState.ERROR.name());
     currentState.setEndTime(_partition.getPartitionName(), System.currentTimeMillis());
     accessor.setProperty(currentStateKey, currentState);

     PropertyKey messageKey =
         new PropertyKey.Builder(_clusterName).message(_initialMaster, messages.get(0).getMsgId());
     accessor.removeProperty(messageKey);

     _fullPipeline.handle(event);

     messages = getMessages(secondMaster);

     Assert.assertEquals(messages.size(), 1);
     Assert.assertEquals(messages.get(0).getTgtName(), secondMaster);
     Assert.assertEquals(messages.get(0).getFromState(), MasterSlaveSMD.States.SLAVE.name());
     Assert.assertEquals(messages.get(0).getToState(), MasterSlaveSMD.States.MASTER.name());
   }

   @Test
   public void testP2PWithInstanceOffline() throws Exception {
     reset(_initialStateMap);
     // Disable old master ((initialMaster) instance,
     // Validate: a M->S message should be sent to initialMaster with a P2P message attached for secondMaster.

     // disable existing master instance
     admin.enableInstance(_clusterName, _initialMaster, false);
     _dataCache.notifyDataChange(HelixConstants.ChangeType.INSTANCE_CONFIG);
     _fullPipeline.handle(event);

     ResourcesStateMap bestpossibleState =
         event.getAttribute(AttributeName.INTERMEDIATE_STATE.name());

     List<Message> messages = getMessages(_initialMaster);

     Assert.assertEquals(messages.size(), 1);
     Message toSlaveMessage = messages.get(0);
     ;

     // verify p2p message are attached to the M->S message sent to the old master instance
     Assert.assertEquals(toSlaveMessage.getRelayMessages().entrySet().size(), 1);
     String secondMaster =
         getTopStateInstance(bestpossibleState.getInstanceStateMap(_db, _partition),
             MasterSlaveSMD.States.MASTER.name());
     Message relayMessage = toSlaveMessage.getRelayMessage(secondMaster);
     Assert.assertNotNull(relayMessage);
     Assert.assertEquals(relayMessage.getMsgSubType(), Message.MessageType.RELAYED_MESSAGE.name());
     Assert.assertEquals(relayMessage.getTgtName(), secondMaster);

     // Old master (initialMaster) completed the M->S transition,
     // but has not forward p2p message to new master (secondMaster) yet.
     // Validate: Controller should not send S->M message to new master.
     handleMessage(_initialMaster, _db);
     _fullPipeline.handle(event);
     messages = getMessages(secondMaster);
     Assert.assertEquals(messages.size(), 0);

     // New master (second master) instance goes offline, controller should send S->M to the third master immediately.
     PropertyKey liveInstanceKey = new PropertyKey.Builder(_clusterName).liveInstance(secondMaster);
     accessor.removeProperty(liveInstanceKey);
     _dataCache.requireFullRefresh();

     _fullPipeline.handle(event);

     bestpossibleState = event.getAttribute(AttributeName.INTERMEDIATE_STATE.name());
     String thirdMaster = getTopStateInstance(bestpossibleState.getInstanceStateMap(_db, _partition),
         MasterSlaveSMD.States.MASTER.name());

     Assert.assertTrue(secondMaster != thirdMaster);
     messages = getMessages(thirdMaster);
     Assert.assertEquals(messages.size(), 1);
     Assert.assertEquals(messages.get(0).getTgtName(), thirdMaster);
     Assert.assertEquals(messages.get(0).getFromState(), MasterSlaveSMD.States.SLAVE.name());
     Assert.assertEquals(messages.get(0).getToState(), MasterSlaveSMD.States.MASTER.name());
   }

   /**
    * This is to simulate the participant (without starting a real participant thread) to handle the pending message.
    * It sets the CurrentState to target State, and remove the pending message from ZK.
    * @param instance
    * @param resource
    */
   private void handleMessage(String instance, String resource) {
     PropertyKey propertyKey = new PropertyKey.Builder(_clusterName).messages(instance);
     List<Message> messages = accessor.getChildValues(propertyKey, true);
     String session = _dataCache.getLiveInstances().get(instance).getEphemeralOwner();

     for (Message m : messages) {
       if (m.getResourceName().equals(resource)) {
         PropertyKey currentStateKey =
             new PropertyKey.Builder(_clusterName).currentState(instance, session, resource);
         CurrentState currentState = accessor.getProperty(currentStateKey);
         if (currentState == null) {
           currentState = new CurrentState(resource);
           currentState.setSessionId(session);
           currentState.setStateModelDefRef(BuiltInStateModelDefinitions.MasterSlave.name());
         }

         String partition = m.getPartitionName();
         String fromState = m.getFromState();
         String toState = m.getToState();
         String partitionState = currentState.getState(partition);

         if ((partitionState == null && fromState.equals(
             BuiltInStateModelDefinitions.MasterSlave.getStateModelDefinition().getInitialState()))
             || (partitionState.equals(fromState))) {
           currentState.setPreviousState(partition, fromState);
           currentState.setState(partition, toState);
           currentState.setStartTime(partition, System.currentTimeMillis());
           try {
             Thread.sleep(50);
           } catch (InterruptedException e) {
             e.printStackTrace();
           }
           currentState.setEndTime(partition, System.currentTimeMillis());
           accessor.setProperty(currentStateKey, currentState);
           PropertyKey messageKey =
               new PropertyKey.Builder(_clusterName).message(instance, m.getMsgId());
           accessor.removeProperty(messageKey);
         }
       }
     }
   }

   /**
    *  Enable all instances, clean all pending messages, set CurrentState to the BestPossibleState
     */
   private void reset(ResourcesStateMap bestpossibleState) {
     for (String ins : _liveInstanceMap.keySet()) {
       LiveInstance liveInstance = _liveInstanceMap.get(ins);
       PropertyKey.Builder keyBuilder = accessor.keyBuilder();
       accessor.setProperty(keyBuilder.liveInstance(liveInstance.getId()), liveInstance);
     }
     for (String ins : _instances) {
       admin.enableInstance(_clusterName, _initialMaster, true);
       cleanMessages(ins);
     }
     for (String resource : bestpossibleState.resourceSet()) {
       setCurrentState(resource, bestpossibleState.getPartitionStateMap(resource).getStateMap());
     }
     for (String ins : _instances) {
       cleanMessages(ins);
     }
     _dataCache.requireFullRefresh();
   }

   private void setCurrentState(String resource,
       Map<Partition, Map<String, String>> partitionStateMap) {
     for (Partition p : partitionStateMap.keySet()) {
       Map<String, String> partitionState = partitionStateMap.get(p);
       for (String instance : partitionState.keySet()) {
         String state = partitionState.get(instance);
         String session = _liveInstanceMap.get(instance).getEphemeralOwner();
         PropertyKey currentStateKey =
             new PropertyKey.Builder(_clusterName).currentState(instance, session, resource);
         CurrentState currentState = accessor.getProperty(currentStateKey);
         if (currentState == null) {
           currentState = new CurrentState(resource);
           currentState.setSessionId(session);
           currentState.setStateModelDefRef(BuiltInStateModelDefinitions.MasterSlave.name());
         }
         currentState.setState(p.getPartitionName(), state);
         accessor.setProperty(currentStateKey, currentState);
       }
     }
   }

   private void cleanMessages(String instance) {
     PropertyKey propertyKey = new PropertyKey.Builder(_clusterName).messages(instance);
     List<Message> messages = accessor.getChildValues(propertyKey, true);
     for (Message m : messages) {
       accessor
           .removeProperty(new PropertyKey.Builder(_clusterName).message(instance, m.getMsgId()));
     }
   }

   List<Message> getMessages(String instance) {
     return accessor.getChildValues(new PropertyKey.Builder(_clusterName).messages(instance), true);
   }

   private String getTopStateInstance(Map<String, String> instanceStateMap, String topState) {
     String masterInstance = null;
     for (Map.Entry<String, String> e : instanceStateMap.entrySet()) {
       if (topState.equals(e.getValue())) {
         masterInstance = e.getKey();
       }
     }

     return masterInstance;
   }
 }
	package org.apache.helix.messaging.p2pMessage;

	/*
	* 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.lang.reflect.Method;
	import java.util.Date;
	import java.util.List;
	import java.util.Map;
	import java.util.Set;
	import java.util.concurrent.Executors;

	import org.apache.helix.HelixConstants;
	import org.apache.helix.HelixDefinedState;
	import org.apache.helix.PropertyKey;
	import org.apache.helix.controller.common.ResourcesStateMap;
	import org.apache.helix.controller.dataproviders.ResourceControllerDataProvider;
	import org.apache.helix.controller.pipeline.Pipeline;
	import org.apache.helix.controller.stages.AttributeName;
	import org.apache.helix.controller.stages.BaseStageTest;
	import org.apache.helix.controller.stages.BestPossibleStateCalcStage;
	import org.apache.helix.controller.stages.CurrentStateComputationStage;
	import org.apache.helix.controller.stages.IntermediateStateCalcStage;
	import org.apache.helix.controller.stages.MessageGenerationPhase;
	import org.apache.helix.controller.stages.MessageSelectionStage;
	import org.apache.helix.controller.stages.MessageThrottleStage;
	import org.apache.helix.controller.stages.ReadClusterDataStage;
	import org.apache.helix.controller.stages.ResourceComputationStage;
	import org.apache.helix.controller.stages.resource.ResourceMessageDispatchStage;
	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.MasterSlaveSMD;
	import org.apache.helix.model.Message;
	import org.apache.helix.model.Partition;
	import org.testng.Assert;
	import org.testng.ITestContext;
	import org.testng.annotations.BeforeClass;
	import org.testng.annotations.BeforeMethod;
	import org.testng.annotations.Test;

	public class TestP2PMessages extends BaseStageTest {
	private String _db = "testDB";
	private int _numPartition = 1;
	private int _numReplica = 3;

	private Partition _partition = new Partition(_db + "_0");

	private ResourceControllerDataProvider _dataCache;
	private Pipeline _fullPipeline;

	private ResourcesStateMap _initialStateMap;

	private Set<String> _instances;
	private Map<String, LiveInstance> _liveInstanceMap;
	private String _initialMaster;

	@BeforeClass
	public void beforeClass() {
	super.beforeClass();
	setup();
	setupIdealState(3, new String[] { _db }, _numPartition, _numReplica,
	IdealState.RebalanceMode.SEMI_AUTO, BuiltInStateModelDefinitions.MasterSlave.name());
	setupStateModel();
	setupInstances(3);
	setupLiveInstances(3);

	ClusterConfig clusterConfig = new ClusterConfig(_clusterName);
	clusterConfig.enableP2PMessage(true);
	setClusterConfig(clusterConfig);

	_dataCache = new ResourceControllerDataProvider(_clusterName);
	_dataCache.setAsyncTasksThreadPool(Executors.newSingleThreadExecutor());

	_dataCache.refresh(manager.getHelixDataAccessor());

	event.addAttribute(AttributeName.ControllerDataProvider.name(), _dataCache);
	event.addAttribute(AttributeName.helixmanager.name(), manager);

	_fullPipeline = new Pipeline("FullPipeline");
	_fullPipeline.addStage(new ReadClusterDataStage());
	_fullPipeline.addStage(new ResourceComputationStage());
	_fullPipeline.addStage(new CurrentStateComputationStage());
	_fullPipeline.addStage(new BestPossibleStateCalcStage());
	_fullPipeline.addStage(new MessageGenerationPhase());
	_fullPipeline.addStage(new MessageSelectionStage());
	_fullPipeline.addStage(new IntermediateStateCalcStage());
	_fullPipeline.addStage(new MessageThrottleStage());
	_fullPipeline.addStage(new ResourceMessageDispatchStage());

	try {
	_fullPipeline.handle(event);
	} catch (Exception e) {
	e.printStackTrace();
	}

	_instances = _dataCache.getAllInstances();
	_liveInstanceMap = _dataCache.getLiveInstances();

	_initialStateMap = event.getAttribute(AttributeName.BEST_POSSIBLE_STATE.name());
	_initialMaster = getTopStateInstance(_initialStateMap.getInstanceStateMap(_db, _partition),
	MasterSlaveSMD.States.MASTER.name());
	Assert.assertNotNull(_initialMaster);
	}

	@BeforeMethod // just to overide the per-test setup in base class.
	public void beforeTest(Method testMethod, ITestContext testContext) {
	long startTime = System.currentTimeMillis();
	System.out.println("START " + testMethod.getName() + " at " + new Date(startTime));
	testContext.setAttribute("StartTime", System.currentTimeMillis());
	}

	@Test
	public void testP2PSendAndTimeout() throws Exception {
	reset(_initialStateMap);

	// Disable old master ((initialMaster) instance,
	// Validate: a M->S message should be sent to initialMaster with a P2P message attached for secondMaster.
	admin.enableInstance(_clusterName, _initialMaster, false);
	_dataCache = event.getAttribute(AttributeName.ControllerDataProvider.name());
	_dataCache.notifyDataChange(HelixConstants.ChangeType.INSTANCE_CONFIG);

	_fullPipeline.handle(event);

	ResourcesStateMap bestpossibleState =
	event.getAttribute(AttributeName.INTERMEDIATE_STATE.name());

	List<Message> messages = getMessages(_initialMaster);

	Assert.assertEquals(messages.size(), 1);
	Message toSlaveMessage = messages.get(0);
	Assert.assertEquals(toSlaveMessage.getTgtName(), _initialMaster);
	Assert.assertEquals(toSlaveMessage.getFromState(), MasterSlaveSMD.States.MASTER.name());
	Assert.assertEquals(toSlaveMessage.getToState(), MasterSlaveSMD.States.SLAVE.name());

	// verify p2p message are attached to the M->S message sent to the old master instance
	Assert.assertEquals(toSlaveMessage.getRelayMessages().entrySet().size(), 1);
	String secondMaster =
	getTopStateInstance(bestpossibleState.getInstanceStateMap(_db, _partition),
	MasterSlaveSMD.States.MASTER.name());

	Message relayMessage = toSlaveMessage.getRelayMessage(secondMaster);
	Assert.assertNotNull(relayMessage);
	Assert.assertEquals(relayMessage.getMsgSubType(), Message.MessageType.RELAYED_MESSAGE.name());
	Assert.assertEquals(relayMessage.getTgtName(), secondMaster);
	Assert.assertEquals(relayMessage.getRelaySrcHost(), _initialMaster);
	Assert.assertEquals(relayMessage.getFromState(), MasterSlaveSMD.States.SLAVE.name());
	Assert.assertEquals(relayMessage.getToState(), MasterSlaveSMD.States.MASTER.name());

	// Old master (initialMaster) completed the M->S transition,
	// but has not forward p2p message to new master (secondMaster) yet.
	// Validate: Controller should not send S->M message to new master.
	handleMessage(_initialMaster, _db);
	_fullPipeline.handle(event);
	messages = getMessages(secondMaster);
	Assert.assertEquals(messages.size(), 0);

	// Old master (initialMaster) completed the M->S transition,
	// but has not forward p2p message to new master (secondMaster) yet, but p2p message should already timeout.
	// Validate: Controller should send S->M message to new master.
	Thread.sleep(Message.RELAY_MESSAGE_DEFAULT_EXPIRY);

	_fullPipeline.handle(event);
	messages = getMessages(secondMaster);
	Assert.assertEquals(messages.size(), 1);
	Assert.assertEquals(messages.get(0).getTgtName(), secondMaster);
	Assert.assertEquals(messages.get(0).getFromState(), MasterSlaveSMD.States.SLAVE.name());
	Assert.assertEquals(messages.get(0).getToState(), MasterSlaveSMD.States.MASTER.name());
	}

	@Test
	public void testP2PWithErrorState() throws Exception {
	reset(_initialStateMap);
	// Disable old master ((initialMaster) instance,
	// Validate: a M->S message should be sent to initialMaster with a P2P message attached for secondMaster.

	// disable existing master instance
	admin.enableInstance(_clusterName, _initialMaster, false);
	_dataCache.notifyDataChange(HelixConstants.ChangeType.INSTANCE_CONFIG);
	_fullPipeline.handle(event);

	ResourcesStateMap bestpossibleState =
	event.getAttribute(AttributeName.INTERMEDIATE_STATE.name());

	List<Message> messages = getMessages(_initialMaster);

	Assert.assertEquals(messages.size(), 1);
	Message toSlaveMessage = messages.get(0);

	// verify p2p message are attached to the M->S message sent to the old master instance
	Assert.assertEquals(toSlaveMessage.getRelayMessages().entrySet().size(), 1);
	String secondMaster =
	getTopStateInstance(bestpossibleState.getInstanceStateMap(_db, _partition),
	MasterSlaveSMD.States.MASTER.name());
	Message relayMessage = toSlaveMessage.getRelayMessage(secondMaster);
	Assert.assertNotNull(relayMessage);
	Assert.assertEquals(relayMessage.getMsgSubType(), Message.MessageType.RELAYED_MESSAGE.name());
	Assert.assertEquals(relayMessage.getTgtName(), secondMaster);

	// Old master (initialMaster) failed the M->S transition,
	// but has not forward p2p message to new master (secondMaster) yet.
	// Validate: Controller should ignore the ERROR partition and send S->M message to new master.
	String session = _dataCache.getLiveInstances().get(_initialMaster).getEphemeralOwner();
	PropertyKey currentStateKey =
	new PropertyKey.Builder(_clusterName).currentState(_initialMaster, session, _db);
	CurrentState currentState = accessor.getProperty(currentStateKey);
	currentState
	.setPreviousState(_partition.getPartitionName(), MasterSlaveSMD.States.MASTER.name());
	currentState.setState(_partition.getPartitionName(), HelixDefinedState.ERROR.name());
	currentState.setEndTime(_partition.getPartitionName(), System.currentTimeMillis());
	accessor.setProperty(currentStateKey, currentState);

	PropertyKey messageKey =
	new PropertyKey.Builder(_clusterName).message(_initialMaster, messages.get(0).getMsgId());
	accessor.removeProperty(messageKey);

	_fullPipeline.handle(event);

	messages = getMessages(secondMaster);

	Assert.assertEquals(messages.size(), 1);
	Assert.assertEquals(messages.get(0).getTgtName(), secondMaster);
	Assert.assertEquals(messages.get(0).getFromState(), MasterSlaveSMD.States.SLAVE.name());
	Assert.assertEquals(messages.get(0).getToState(), MasterSlaveSMD.States.MASTER.name());
	}

	@Test
	public void testP2PWithInstanceOffline() throws Exception {
	reset(_initialStateMap);
	// Disable old master ((initialMaster) instance,
	// Validate: a M->S message should be sent to initialMaster with a P2P message attached for secondMaster.

	// disable existing master instance
	admin.enableInstance(_clusterName, _initialMaster, false);
	_dataCache.notifyDataChange(HelixConstants.ChangeType.INSTANCE_CONFIG);
	_fullPipeline.handle(event);

	ResourcesStateMap bestpossibleState =
	event.getAttribute(AttributeName.INTERMEDIATE_STATE.name());

	List<Message> messages = getMessages(_initialMaster);

	Assert.assertEquals(messages.size(), 1);
	Message toSlaveMessage = messages.get(0);
	;

	// verify p2p message are attached to the M->S message sent to the old master instance
	Assert.assertEquals(toSlaveMessage.getRelayMessages().entrySet().size(), 1);
	String secondMaster =
	getTopStateInstance(bestpossibleState.getInstanceStateMap(_db, _partition),
	MasterSlaveSMD.States.MASTER.name());
	Message relayMessage = toSlaveMessage.getRelayMessage(secondMaster);
	Assert.assertNotNull(relayMessage);
	Assert.assertEquals(relayMessage.getMsgSubType(), Message.MessageType.RELAYED_MESSAGE.name());
	Assert.assertEquals(relayMessage.getTgtName(), secondMaster);

	// Old master (initialMaster) completed the M->S transition,
	// but has not forward p2p message to new master (secondMaster) yet.
	// Validate: Controller should not send S->M message to new master.
	handleMessage(_initialMaster, _db);
	_fullPipeline.handle(event);
	messages = getMessages(secondMaster);
	Assert.assertEquals(messages.size(), 0);

	// New master (second master) instance goes offline, controller should send S->M to the third master immediately.
	PropertyKey liveInstanceKey = new PropertyKey.Builder(_clusterName).liveInstance(secondMaster);
	accessor.removeProperty(liveInstanceKey);
	_dataCache.requireFullRefresh();

	_fullPipeline.handle(event);

	bestpossibleState = event.getAttribute(AttributeName.INTERMEDIATE_STATE.name());
	String thirdMaster = getTopStateInstance(bestpossibleState.getInstanceStateMap(_db, _partition),
	MasterSlaveSMD.States.MASTER.name());

	Assert.assertTrue(secondMaster != thirdMaster);
	messages = getMessages(thirdMaster);
	Assert.assertEquals(messages.size(), 1);
	Assert.assertEquals(messages.get(0).getTgtName(), thirdMaster);
	Assert.assertEquals(messages.get(0).getFromState(), MasterSlaveSMD.States.SLAVE.name());
	Assert.assertEquals(messages.get(0).getToState(), MasterSlaveSMD.States.MASTER.name());
	}

	/**
	* This is to simulate the participant (without starting a real participant thread) to handle the pending message.
	* It sets the CurrentState to target State, and remove the pending message from ZK.
	* @param instance
	* @param resource
	*/
	private void handleMessage(String instance, String resource) {
	PropertyKey propertyKey = new PropertyKey.Builder(_clusterName).messages(instance);
	List<Message> messages = accessor.getChildValues(propertyKey, true);
	String session = _dataCache.getLiveInstances().get(instance).getEphemeralOwner();

	for (Message m : messages) {
	if (m.getResourceName().equals(resource)) {
	PropertyKey currentStateKey =
	new PropertyKey.Builder(_clusterName).currentState(instance, session, resource);
	CurrentState currentState = accessor.getProperty(currentStateKey);
	if (currentState == null) {
	currentState = new CurrentState(resource);
	currentState.setSessionId(session);
	currentState.setStateModelDefRef(BuiltInStateModelDefinitions.MasterSlave.name());
	}

	String partition = m.getPartitionName();
	String fromState = m.getFromState();
	String toState = m.getToState();
	String partitionState = currentState.getState(partition);

	if ((partitionState == null && fromState.equals(
	BuiltInStateModelDefinitions.MasterSlave.getStateModelDefinition().getInitialState()))
	\|\| (partitionState.equals(fromState))) {
	currentState.setPreviousState(partition, fromState);
	currentState.setState(partition, toState);
	currentState.setStartTime(partition, System.currentTimeMillis());
	try {
	Thread.sleep(50);
	} catch (InterruptedException e) {
	e.printStackTrace();
	}
	currentState.setEndTime(partition, System.currentTimeMillis());
	accessor.setProperty(currentStateKey, currentState);
	PropertyKey messageKey =
	new PropertyKey.Builder(_clusterName).message(instance, m.getMsgId());
	accessor.removeProperty(messageKey);
	}
	}
	}
	}

	/**
	* Enable all instances, clean all pending messages, set CurrentState to the BestPossibleState
	*/
	private void reset(ResourcesStateMap bestpossibleState) {
	for (String ins : _liveInstanceMap.keySet()) {
	LiveInstance liveInstance = _liveInstanceMap.get(ins);
	PropertyKey.Builder keyBuilder = accessor.keyBuilder();
	accessor.setProperty(keyBuilder.liveInstance(liveInstance.getId()), liveInstance);
	}
	for (String ins : _instances) {
	admin.enableInstance(_clusterName, _initialMaster, true);
	cleanMessages(ins);
	}
	for (String resource : bestpossibleState.resourceSet()) {
	setCurrentState(resource, bestpossibleState.getPartitionStateMap(resource).getStateMap());
	}
	for (String ins : _instances) {
	cleanMessages(ins);
	}
	_dataCache.requireFullRefresh();
	}

	private void setCurrentState(String resource,
	Map<Partition, Map<String, String>> partitionStateMap) {
	for (Partition p : partitionStateMap.keySet()) {
	Map<String, String> partitionState = partitionStateMap.get(p);
	for (String instance : partitionState.keySet()) {
	String state = partitionState.get(instance);
	String session = _liveInstanceMap.get(instance).getEphemeralOwner();
	PropertyKey currentStateKey =
	new PropertyKey.Builder(_clusterName).currentState(instance, session, resource);
	CurrentState currentState = accessor.getProperty(currentStateKey);
	if (currentState == null) {
	currentState = new CurrentState(resource);
	currentState.setSessionId(session);
	currentState.setStateModelDefRef(BuiltInStateModelDefinitions.MasterSlave.name());
	}
	currentState.setState(p.getPartitionName(), state);
	accessor.setProperty(currentStateKey, currentState);
	}
	}
	}

	private void cleanMessages(String instance) {
	PropertyKey propertyKey = new PropertyKey.Builder(_clusterName).messages(instance);
	List<Message> messages = accessor.getChildValues(propertyKey, true);
	for (Message m : messages) {
	accessor
	.removeProperty(new PropertyKey.Builder(_clusterName).message(instance, m.getMsgId()));
	}
	}

	List<Message> getMessages(String instance) {
	return accessor.getChildValues(new PropertyKey.Builder(_clusterName).messages(instance), true);
	}

	private String getTopStateInstance(Map<String, String> instanceStateMap, String topState) {
	String masterInstance = null;
	for (Map.Entry<String, String> e : instanceStateMap.entrySet()) {
	if (topState.equals(e.getValue())) {
	masterInstance = e.getKey();
	}
	}

	return masterInstance;
	}
	}