Java Code Examples for org.apache.helix.model.Message#setTgtSessionId()

private Message generateMessage(String from, String to) {
  String uuid = UUID.randomUUID().toString();
  Message message = new Message(Message.MessageType.STATE_TRANSITION, uuid);
  message.setSrcName("ADMIN");
  message.setTgtName(_participant.getInstanceName());
  message.setMsgState(Message.MessageState.NEW);
  message.setPartitionName("P");
  message.setResourceName(WorkflowGenerator.DEFAULT_TGT_DB);
  message.setFromState(from);
  message.setToState(to);
  message.setTgtSessionId(_participant.getSessionId());
  message.setSrcSessionId(_manager.getSessionId());
  message.setStateModelDef("OnlineOffline");
  message.setStateModelFactoryName("DEFAULT");
  return message;
}
 

private void syncSessionToController(HelixManager manager) {
  if (_lastSessionSyncTime == null ||
          System.currentTimeMillis() - _lastSessionSyncTime > SESSION_SYNC_INTERVAL) { // > delay since last sync
    HelixDataAccessor accessor = manager.getHelixDataAccessor();
    PropertyKey key = new Builder(manager.getClusterName()).controllerMessage(SESSION_SYNC);
    if (accessor.getProperty(key) == null) {
      LOG.info(String.format("Participant %s syncs session with controller", manager.getInstanceName()));
      Message msg = new Message(MessageType.PARTICIPANT_SESSION_CHANGE, SESSION_SYNC);
      msg.setSrcName(manager.getInstanceName());
      msg.setTgtSessionId("*");
      msg.setMsgState(MessageState.NEW);
      msg.setMsgId(SESSION_SYNC);

      Criteria cr = new Criteria();
      cr.setRecipientInstanceType(InstanceType.CONTROLLER);
      cr.setSessionSpecific(false);

      manager.getMessagingService().send(cr, msg);
      _lastSessionSyncTime = System.currentTimeMillis();
    }
  }
}
 

void sendMessage(String msgId, String instanceName, String fromState, String toState,
    String partitionKey, int partitionId) throws InterruptedException, JsonGenerationException,
    JsonMappingException, IOException {
  Message message = new Message(MessageType.STATE_TRANSITION, msgId);
  message.setMsgId(msgId);
  message.setSrcName(srcName);
  message.setTgtName(instanceName);
  message.setMsgState(MessageState.NEW);
  message.setFromState(fromState);
  message.setToState(toState);
  // message.setPartitionId(partitionId);
  message.setPartitionName(partitionKey);

  String path = PropertyPathBuilder.instanceMessage(clusterName, instanceName, message.getId());
  ObjectMapper mapper = new ObjectMapper();
  StringWriter sw = new StringWriter();
  mapper.writeValueUsingView(sw, message, Message.class);
  System.out.println(sw.toString());
  client.delete(path);

  Thread.sleep(10000);
  ZNRecord record = client.readData(PropertyPathBuilder.liveInstance(clusterName, instanceName));
  message.setTgtSessionId(record.getSimpleField(LiveInstanceProperty.SESSION_ID.toString())
      .toString());
  client.createPersistent(path, message);
}
 

private Message newMessage(String resourceName, String partitionName, String instanceName,
    String fromState, String toState) {
  String uuid = UUID.randomUUID().toString();
  Message message = new Message(MessageType.STATE_TRANSITION, uuid);
  message.setSrcName("controller");
  message.setTgtName(instanceName);
  message.setMsgState(MessageState.NEW);
  message.setResourceName(resourceName);
  message.setPartitionName(partitionName);
  message.setFromState(fromState);
  message.setToState(toState);
  message.setTgtSessionId("sessionId");
  message.setSrcSessionId("sessionId");
  message.setStateModelDef("MasterSlave");
  message.setStateModelFactoryName("DEFAULT");
  message.setBucketSize(0);
  return message;
}
 

@VisibleForTesting
void sendShutdownRequest() {
  Criteria criteria = new Criteria();
  criteria.setInstanceName("%");
  criteria.setPartition("%");
  criteria.setPartitionState("%");
  criteria.setResource("%");
  if (this.isHelixClusterManaged) {
    //In the managed mode, the Gobblin Yarn Application Master connects to the Helix cluster in the Participant role.
    criteria.setRecipientInstanceType(InstanceType.PARTICIPANT);
    criteria.setInstanceName(this.helixInstanceName);
  } else {
    criteria.setRecipientInstanceType(InstanceType.CONTROLLER);
  }
  criteria.setSessionSpecific(true);

  Message shutdownRequest = new Message(GobblinHelixConstants.SHUTDOWN_MESSAGE_TYPE,
      HelixMessageSubTypes.APPLICATION_MASTER_SHUTDOWN.toString().toLowerCase() + UUID.randomUUID().toString());
  shutdownRequest.setMsgSubType(HelixMessageSubTypes.APPLICATION_MASTER_SHUTDOWN.toString());
  shutdownRequest.setMsgState(Message.MessageState.NEW);
  shutdownRequest.setTgtSessionId("*");

  int messagesSent = this.messagingService.send(criteria, shutdownRequest);
  if (messagesSent == 0) {
    LOGGER.error(String.format("Failed to send the %s message to the controller", shutdownRequest.getMsgSubType()));
  }
}
 

@VisibleForTesting
protected void sendTokenFileUpdatedMessage(InstanceType instanceType, String instanceName) {
  Criteria criteria = new Criteria();
  criteria.setInstanceName(Strings.isNullOrEmpty(instanceName) ? "%" : instanceName);
  criteria.setResource("%");
  criteria.setPartition("%");
  criteria.setPartitionState("%");
  criteria.setRecipientInstanceType(instanceType);
  /**
   * #HELIX-0.6.7-WORKAROUND
   * Add back when LIVESTANCES messaging is ported to 0.6 branch
   if (instanceType == InstanceType.PARTICIPANT) {
   criteria.setDataSource(Criteria.DataSource.LIVEINSTANCES);
   }
   **/
  criteria.setSessionSpecific(true);

  Message tokenFileUpdatedMessage = new Message(Message.MessageType.USER_DEFINE_MSG,
      HelixMessageSubTypes.TOKEN_FILE_UPDATED.toString().toLowerCase() + UUID.randomUUID().toString());
  tokenFileUpdatedMessage.setMsgSubType(HelixMessageSubTypes.TOKEN_FILE_UPDATED.toString());
  tokenFileUpdatedMessage.setMsgState(Message.MessageState.NEW);
  if (instanceType == InstanceType.CONTROLLER) {
    tokenFileUpdatedMessage.setTgtSessionId("*");
  }

  // #HELIX-0.6.7-WORKAROUND
  // Temporarily bypass the default messaging service to allow upgrade to 0.6.7 which is missing support
  // for messaging to instances
  //int messagesSent = this.helixManager.getMessagingService().send(criteria, tokenFileUpdatedMessage);
  GobblinHelixMessagingService messagingService = new GobblinHelixMessagingService(helixManager);

  int messagesSent = messagingService.send(criteria, tokenFileUpdatedMessage);
  LOGGER.info(String.format("Sent %d token file updated message(s) to the %s", messagesSent, instanceType));
}
 

private Message createStateTransitionCancellationMessage(HelixManager manager, Resource resource,
    String partitionName, String instanceName, String sessionId, String stateModelDefName,
    String fromState, String toState, String nextState, Message cancellationMessage,
    boolean isCancellationEnabled, String currentState) {

  if (isCancellationEnabled && cancellationMessage == null) {
    LogUtil.logInfo(logger, _eventId,
        "Send cancellation message of the state transition for " + resource.getResourceName()
            + "." + partitionName + " on " + instanceName + ", currentState: " + currentState
            + ", nextState: " + (nextState == null ? "N/A" : nextState));

    String uuid = UUID.randomUUID().toString();
    Message message = new Message(MessageType.STATE_TRANSITION_CANCELLATION, uuid);
    message.setSrcName(manager.getInstanceName());
    message.setTgtName(instanceName);
    message.setMsgState(MessageState.NEW);
    message.setPartitionName(partitionName);
    message.setResourceName(resource.getResourceName());
    message.setFromState(fromState);
    message.setToState(toState);
    message.setTgtSessionId(sessionId);
    message.setSrcSessionId(manager.getSessionId());
    message.setStateModelDef(stateModelDefName);
    message.setStateModelFactoryName(resource.getStateModelFactoryname());
    message.setBucketSize(resource.getBucketSize());
    return message;
  }

  return null;
}
 

private List<Message> generateMessagesForParticipant(Criteria recipientCriteria, Message message,
    HelixDataAccessor targetDataAccessor) {
  List<Message> messages = new ArrayList<Message>();
  List<Map<String, String>> matchedList =
      _evaluator.evaluateCriteria(recipientCriteria, targetDataAccessor);

  if (!matchedList.isEmpty()) {
    Map<String, String> sessionIdMap = new HashMap<String, String>();
    if (recipientCriteria.isSessionSpecific()) {
      Builder keyBuilder = targetDataAccessor.keyBuilder();
      // For backward compatibility, allow partial read for the live instances.
      // Note that this may cause the pending message to be sent with null target session Id.
      List<LiveInstance> liveInstances =
          targetDataAccessor.getChildValues(keyBuilder.liveInstances(), false);

      for (LiveInstance liveInstance : liveInstances) {
        sessionIdMap.put(liveInstance.getInstanceName(), liveInstance.getEphemeralOwner());
      }
    }
    for (Map<String, String> map : matchedList) {
      String id = UUID.randomUUID().toString();
      Message newMessage = new Message(message.getRecord(), id);
      String srcInstanceName = _manager.getInstanceName();
      String tgtInstanceName = map.get("instanceName");
      // Don't send message to self
      if (recipientCriteria.isSelfExcluded() && srcInstanceName.equalsIgnoreCase(tgtInstanceName)) {
        continue;
      }
      newMessage.setSrcName(srcInstanceName);
      newMessage.setTgtName(tgtInstanceName);
      newMessage.setResourceName(map.get("resourceName"));
      newMessage.setPartitionName(map.get("partitionName"));
      if (recipientCriteria.isSessionSpecific()) {
        newMessage.setTgtSessionId(sessionIdMap.get(tgtInstanceName));
      }
      messages.add(newMessage);
    }
  }
  return messages;
}
 

@Test()
public void testInvocation() throws Exception {
  HelixTaskExecutor executor = new HelixTaskExecutor();
  System.out.println("START TestCMTaskHandler.testInvocation()");
  Message message = new Message(MessageType.STATE_TRANSITION, "Some unique id");

  message.setSrcName("cm-instance-0");
  message.setTgtSessionId("1234");
  message.setFromState("Offline");
  message.setToState("Slave");
  message.setPartitionName("TestDB_0");
  message.setMsgId("Some unique message id");
  message.setResourceName("TestDB");
  message.setTgtName("localhost");
  message.setStateModelDef("MasterSlave");
  message.setStateModelFactoryName(HelixConstants.DEFAULT_STATE_MODEL_FACTORY);
  MockMasterSlaveStateModel stateModel = new MockMasterSlaveStateModel();
  NotificationContext context;
  MockManager manager = new MockManager("clusterName");
  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  StateModelDefinition stateModelDef =
      new StateModelDefinition(StateModelConfigGenerator.generateConfigForMasterSlave());
  Builder keyBuilder = accessor.keyBuilder();
  accessor.setProperty(keyBuilder.stateModelDef("MasterSlave"), stateModelDef);

  context = new NotificationContext(manager);
  CurrentState currentStateDelta = new CurrentState("TestDB");
  currentStateDelta.setState("TestDB_0", "OFFLINE");

  HelixStateTransitionHandler stHandler =
      new HelixStateTransitionHandler(null, stateModel, message, context, currentStateDelta);
  HelixTask handler;
  handler = new HelixTask(message, context, stHandler, executor);
  handler.call();
  AssertJUnit.assertTrue(stateModel.stateModelInvoked);
  System.out.println("END TestCMTaskHandler.testInvocation() at "
      + new Date(System.currentTimeMillis()));
}
 

@Test()
public void TestBlockingSendReceive() throws Exception {
  String hostSrc = "localhost_" + START_PORT;
  String hostDest = "localhost_" + (START_PORT + 1);

  TestMessagingHandlerFactory factory = new TestMessagingHandlerFactory();
  _participants[1].getMessagingService().registerMessageHandlerFactory(factory.getMessageTypes(),
      factory);

  String msgId = new UUID(123, 456).toString();
  Message msg = new Message(factory.getMessageTypes().get(0), msgId);
  msg.setMsgId(msgId);
  msg.setSrcName(hostSrc);

  msg.setTgtSessionId("*");
  msg.setMsgState(MessageState.NEW);
  String para = "Testing messaging para";
  msg.getRecord().setSimpleField("TestMessagingPara", para);

  Criteria cr = new Criteria();
  cr.setInstanceName(hostDest);
  cr.setRecipientInstanceType(InstanceType.PARTICIPANT);
  cr.setSessionSpecific(false);

  AsyncCallback asyncCallback = new MockAsyncCallback();
  int messagesSent =
      _participants[0].getMessagingService().sendAndWait(cr, msg, asyncCallback, 60000);

  AssertJUnit.assertTrue(asyncCallback.getMessageReplied().get(0).getRecord()
      .getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ReplyMessage")
      .equals("TestReplyMessage"));
  AssertJUnit.assertTrue(asyncCallback.getMessageReplied().size() == 1);

  AsyncCallback asyncCallback2 = new MockAsyncCallback();
  messagesSent = _participants[0].getMessagingService().sendAndWait(cr, msg, asyncCallback2, 500);
  AssertJUnit.assertTrue(asyncCallback2.isTimedOut());

}
 

@Test()
public void sendSelfMsg() {
  String hostSrc = "localhost_" + START_PORT;

  for (int i = 0; i < NODE_NR; i++) {
    TestMessagingHandlerFactory factory = new TestMessagingHandlerFactory();
    String hostDest = "localhost_" + (START_PORT + i);
    _participants[i].getMessagingService().registerMessageHandlerFactory(
        factory.getMessageTypes(), factory);

  }
  String msgId = new UUID(123, 456).toString();
  Message msg = new Message(new TestMessagingHandlerFactory().getMessageTypes().get(0), msgId);
  msg.setMsgId(msgId);
  msg.setSrcName(hostSrc);

  msg.setTgtSessionId("*");
  msg.setMsgState(MessageState.NEW);
  String para = "Testing messaging para";
  msg.getRecord().setSimpleField("TestMessagingPara", para);

  Criteria cr = new Criteria();
  cr.setInstanceName("%");
  cr.setRecipientInstanceType(InstanceType.PARTICIPANT);
  cr.setSessionSpecific(false);
  cr.setSelfExcluded(false);
  AsyncCallback callback1 = new MockAsyncCallback();
  int messageSent1 =
      _participants[0].getMessagingService().sendAndWait(cr, msg, callback1, 10000);

  AssertJUnit.assertTrue(callback1.getMessageReplied().size() == NODE_NR);
  AssertJUnit.assertTrue(callback1.getMessageReplied().get(0).getRecord()
      .getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ReplyMessage")
      .equals("TestReplyMessage"));
}
 

@Test()
public void testDuplicatedMessage() throws InterruptedException {
  System.out.println("START TestHelixTaskExecutor.testDuplicatedMessage()");
  HelixTaskExecutor executor = new HelixTaskExecutor();
  HelixManager manager = new MockClusterManager();
  HelixDataAccessor dataAccessor = manager.getHelixDataAccessor();
  PropertyKey.Builder keyBuilder = dataAccessor.keyBuilder();

  TestStateTransitionHandlerFactory stateTransitionFactory =
      new TestStateTransitionHandlerFactory(Message.MessageType.STATE_TRANSITION.name(), 1000);
  executor.registerMessageHandlerFactory(Message.MessageType.STATE_TRANSITION.name(),
      stateTransitionFactory);

  NotificationContext changeContext = new NotificationContext(manager);
  List<Message> msgList = new ArrayList<Message>();

  int nMsgs = 3;
  String instanceName = manager.getInstanceName();
  for (int i = 0; i < nMsgs; i++) {
    Message msg =
        new Message(Message.MessageType.STATE_TRANSITION.name(), UUID.randomUUID().toString());
    msg.setTgtSessionId(manager.getSessionId());
    msg.setCreateTimeStamp((long) i);
    msg.setTgtName("Localhost_1123");
    msg.setSrcName("127.101.1.23_2234");
    msg.setPartitionName("Partition");
    msg.setResourceName("Resource");
    msg.setStateModelDef("DummyMasterSlave");
    msg.setFromState("SLAVE");
    msg.setToState("MASTER");
    dataAccessor.setProperty(msg.getKey(keyBuilder, instanceName), msg);
    msgList.add(msg);
  }

  AssertJUnit
      .assertEquals(dataAccessor.getChildValues(keyBuilder.messages(instanceName), true).size(),
          nMsgs);

  changeContext.setChangeType(HelixConstants.ChangeType.MESSAGE);
  executor.onMessage(instanceName, msgList, changeContext);

  Thread.sleep(200);

  // only 1 message is left over - state transition takes 1sec
  Assert.assertEquals(dataAccessor.getChildValues(keyBuilder.messages(instanceName), true).size(),
      1);

  // While a state transition message is going on, another state transition message for same
  // resource / partition comes in, it should be discarded by message handler

  // Mock accessor is modifying message state in memory so we set it back to NEW
  msgList.get(2).setMsgState(MessageState.NEW);
  dataAccessor.setProperty(msgList.get(2).getKey(keyBuilder, instanceName), msgList.get(2));
  executor.onMessage(instanceName, Arrays.asList(msgList.get(2)), changeContext);
  Thread.sleep(200);
  Assert.assertEquals(dataAccessor.getChildValues(keyBuilder.messages(instanceName), true).size(),
      1);

  Thread.sleep(1000);
  Assert.assertEquals(dataAccessor.getChildValues(keyBuilder.messages(instanceName), true).size(),
      0);
  System.out.println("END TestHelixTaskExecutor.testDuplicatedMessage()");
}
 

@Test()
public void testInvocation() throws Exception {

  HelixManager testHelixManager =
      HelixManagerFactory.getZKHelixManager(clusterName, "localhost_8900",
          InstanceType.PARTICIPANT, ZK_ADDR);
  testHelixManager.connect();

  TestZKCallback test = new TestZKCallback();

  TestZKCallback.TestCallbackListener testListener = test.new TestCallbackListener();

  testHelixManager.addMessageListener(testListener, "localhost_8900");
  testHelixManager.addCurrentStateChangeListener(testListener, "localhost_8900",
      testHelixManager.getSessionId());
  testHelixManager.addCustomizedStateRootChangeListener(testListener, "localhost_8900");
  testHelixManager.addConfigChangeListener(testListener);
  testHelixManager.addIdealStateChangeListener(testListener);
  testHelixManager.addExternalViewChangeListener(testListener);
  testHelixManager.addLiveInstanceChangeListener(testListener);
  // Initial add listener should trigger the first execution of the
  // listener callbacks
  AssertJUnit.assertTrue(testListener.configChangeReceived
      & testListener.currentStateChangeReceived & testListener.externalViewChangeReceived
      & testListener.idealStateChangeReceived & testListener.liveInstanceChangeReceived
      & testListener.messageChangeReceived);

  testListener.Reset();
  HelixDataAccessor accessor = testHelixManager.getHelixDataAccessor();
  Builder keyBuilder = accessor.keyBuilder();

  ExternalView extView = new ExternalView("db-12345");
  accessor.setProperty(keyBuilder.externalView("db-12345"), extView);
  Thread.sleep(100);
  AssertJUnit.assertTrue(testListener.externalViewChangeReceived);
  testListener.Reset();

  CurrentState curState = new CurrentState("db-12345");
  curState.setSessionId("sessionId");
  curState.setStateModelDefRef("StateModelDef");
  accessor.setProperty(keyBuilder.currentState("localhost_8900", testHelixManager.getSessionId(),
      curState.getId()), curState);
  Thread.sleep(100);
  AssertJUnit.assertTrue(testListener.currentStateChangeReceived);
  testListener.Reset();

  IdealState idealState = new IdealState("db-1234");
  idealState.setNumPartitions(400);
  idealState.setReplicas(Integer.toString(2));
  idealState.setStateModelDefRef("StateModeldef");
  accessor.setProperty(keyBuilder.idealStates("db-1234"), idealState);
  Thread.sleep(100);
  AssertJUnit.assertTrue(testListener.idealStateChangeReceived);
  testListener.Reset();

  // dummyRecord = new ZNRecord("db-12345");
  // dataAccessor.setProperty(PropertyType.IDEALSTATES, idealState, "db-12345"
  // );
  // Thread.sleep(100);
  // AssertJUnit.assertTrue(testListener.idealStateChangeReceived);
  // testListener.Reset();

  // dummyRecord = new ZNRecord("localhost:8900");
  // List<ZNRecord> recList = new ArrayList<ZNRecord>();
  // recList.add(dummyRecord);

  testListener.Reset();
  Message message = new Message(MessageType.STATE_TRANSITION, UUID.randomUUID().toString());
  message.setTgtSessionId("*");
  message.setResourceName("testResource");
  message.setPartitionName("testPartitionKey");
  message.setStateModelDef("MasterSlave");
  message.setToState("toState");
  message.setFromState("fromState");
  message.setTgtName("testTarget");
  message.setStateModelFactoryName(HelixConstants.DEFAULT_STATE_MODEL_FACTORY);

  accessor.setProperty(keyBuilder.message("localhost_8900", message.getId()), message);
  Thread.sleep(500);
  AssertJUnit.assertTrue(testListener.messageChangeReceived);

  // dummyRecord = new ZNRecord("localhost_9801");
  LiveInstance liveInstance = new LiveInstance("localhost_9801");
  liveInstance.setSessionId(UUID.randomUUID().toString());
  liveInstance.setHelixVersion(UUID.randomUUID().toString());
  accessor.setProperty(keyBuilder.liveInstance("localhost_9801"), liveInstance);
  Thread.sleep(500);
  AssertJUnit.assertTrue(testListener.liveInstanceChangeReceived);
  testListener.Reset();

  // dataAccessor.setNodeConfigs(recList); Thread.sleep(100);
  // AssertJUnit.assertTrue(testListener.configChangeReceived);
  // testListener.Reset();

  accessor.removeProperty(keyBuilder.liveInstance("localhost_8900"));
  accessor.removeProperty(keyBuilder.liveInstance("localhost_9801"));
}
 

@Test
public void testSchedulerZeroMsg() throws Exception {
  _factory._results.clear();
  HelixManager manager = null;
  for (int i = 0; i < NODE_NR; i++) {
    _participants[i].getMessagingService()
        .registerMessageHandlerFactory(_factory.getMessageTypes(), _factory);

    manager = _participants[i]; // _startCMResultMap.get(hostDest)._manager;
  }

  Message schedulerMessage =
      new Message(MessageType.SCHEDULER_MSG + "", UUID.randomUUID().toString());
  schedulerMessage.setTgtSessionId("*");
  schedulerMessage.setTgtName("CONTROLLER");
  // TODO: change it to "ADMIN" ?
  schedulerMessage.setSrcName("CONTROLLER");

  // Template for the individual message sent to each participant
  Message msg = new Message(_factory.getMessageTypes().get(0), "Template");
  msg.setTgtSessionId("*");
  msg.setMsgState(MessageState.NEW);

  // Criteria to send individual messages
  Criteria cr = new Criteria();
  cr.setInstanceName("localhost_DOESNOTEXIST");
  cr.setRecipientInstanceType(InstanceType.PARTICIPANT);
  cr.setSessionSpecific(false);
  cr.setResource("%");
  cr.setPartition("%");

  ObjectMapper mapper = new ObjectMapper();
  SerializationConfig serializationConfig = mapper.getSerializationConfig();
  serializationConfig.set(SerializationConfig.Feature.INDENT_OUTPUT, true);

  StringWriter sw = new StringWriter();
  mapper.writeValue(sw, cr);

  String crString = sw.toString();

  schedulerMessage.getRecord().setSimpleField("Criteria", crString);
  schedulerMessage.getRecord().setMapField("MessageTemplate", msg.getRecord().getSimpleFields());
  schedulerMessage.getRecord().setSimpleField("TIMEOUT", "-1");

  HelixDataAccessor helixDataAccessor = manager.getHelixDataAccessor();
  Builder keyBuilder = helixDataAccessor.keyBuilder();
  PropertyKey controllerMessageKey = keyBuilder.controllerMessage(schedulerMessage.getMsgId());
  helixDataAccessor.setProperty(controllerMessageKey, schedulerMessage);

  Thread.sleep(3000);

  Assert.assertEquals(0, _factory._results.size());
  PropertyKey controllerTaskStatus = keyBuilder
      .controllerTaskStatus(MessageType.SCHEDULER_MSG.name(), schedulerMessage.getMsgId());
  for (int i = 0; i < 10; i++) {
    StatusUpdate update = helixDataAccessor.getProperty(controllerTaskStatus);
    if (update == null || update.getRecord().getMapField("SentMessageCount") == null) {
      Thread.sleep(1000);
    }
  }
  ZNRecord statusUpdate = helixDataAccessor.getProperty(controllerTaskStatus).getRecord();
  Assert.assertEquals(statusUpdate.getMapField("SentMessageCount").get("MessageCount"), "0");
  int count = 0;
  for (Set<String> val : _factory._results.values()) {
    count += val.size();
  }
  Assert.assertEquals(count, 0);
}
 

@Test()
public void TestControllerMessage() throws Exception {
  String hostSrc = "localhost_" + START_PORT;

  for (int i = 0; i < NODE_NR; i++) {
    TestMessagingHandlerFactory factory = new TestMessagingHandlerFactory();
    String hostDest = "localhost_" + (START_PORT + i);
    _participants[i].getMessagingService().registerMessageHandlerFactory(
        factory.getMessageTypes(), factory);

  }
  String msgId = new UUID(123, 456).toString();
  Message msg = new Message(MessageType.CONTROLLER_MSG, msgId);
  msg.setMsgId(msgId);
  msg.setSrcName(hostSrc);

  msg.setTgtSessionId("*");
  msg.setMsgState(MessageState.NEW);
  String para = "Testing messaging para";
  msg.getRecord().setSimpleField("TestMessagingPara", para);

  Criteria cr = new Criteria();
  cr.setInstanceName("*");
  cr.setRecipientInstanceType(InstanceType.CONTROLLER);
  cr.setSessionSpecific(false);

  AsyncCallback callback1 = new MockAsyncCallback();
  int messagesSent =
      _participants[0].getMessagingService().sendAndWait(cr, msg, callback1, 10000);

  AssertJUnit.assertTrue(callback1.getMessageReplied().get(0).getRecord()
      .getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ControllerResult")
      .indexOf(hostSrc) != -1);
  AssertJUnit.assertTrue(callback1.getMessageReplied().size() == 1);

  msgId = UUID.randomUUID().toString();
  msg.setMsgId(msgId);
  cr.setPartition("TestDB_17");
  AsyncCallback callback2 = new MockAsyncCallback();
  messagesSent = _participants[0].getMessagingService().sendAndWait(cr, msg, callback2, 10000);

  AssertJUnit.assertTrue(callback2.getMessageReplied().get(0).getRecord()
      .getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ControllerResult")
      .indexOf(hostSrc) != -1);

  AssertJUnit.assertTrue(callback2.getMessageReplied().size() == 1);

  msgId = UUID.randomUUID().toString();
  msg.setMsgId(msgId);
  cr.setPartitionState("SLAVE");
  AsyncCallback callback3 = new MockAsyncCallback();
  messagesSent = _participants[0].getMessagingService().sendAndWait(cr, msg, callback3, 10000);
  AssertJUnit.assertTrue(callback3.getMessageReplied().get(0).getRecord()
      .getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ControllerResult")
      .indexOf(hostSrc) != -1);

  AssertJUnit.assertTrue(callback3.getMessageReplied().size() == 1);
}
 

@Test()
public void TestMultiMessageCriteria() throws Exception {
  String hostSrc = "localhost_" + START_PORT;

  for (int i = 0; i < NODE_NR; i++) {
    TestMessagingHandlerFactory factory = new TestMessagingHandlerFactory();
    String hostDest = "localhost_" + (START_PORT + i);
    _participants[i].getMessagingService().registerMessageHandlerFactory(
        factory.getMessageTypes(), factory);

  }
  String msgId = new UUID(123, 456).toString();
  Message msg = new Message(new TestMessagingHandlerFactory().getMessageTypes().get(0), msgId);
  msg.setMsgId(msgId);
  msg.setSrcName(hostSrc);

  msg.setTgtSessionId("*");
  msg.setMsgState(MessageState.NEW);
  String para = "Testing messaging para";
  msg.getRecord().setSimpleField("TestMessagingPara", para);

  Criteria cr = new Criteria();
  cr.setInstanceName("%");
  cr.setRecipientInstanceType(InstanceType.PARTICIPANT);
  cr.setSessionSpecific(false);
  AsyncCallback callback1 = new MockAsyncCallback();
  int messageSent1 =
      _participants[0].getMessagingService().sendAndWait(cr, msg, callback1, 10000);

  AssertJUnit.assertTrue(callback1.getMessageReplied().get(0).getRecord()
      .getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ReplyMessage")
      .equals("TestReplyMessage"));
  AssertJUnit.assertTrue(callback1.getMessageReplied().size() == NODE_NR - 1);

  AsyncCallback callback2 = new MockAsyncCallback();
  int messageSent2 = _participants[0].getMessagingService().sendAndWait(cr, msg, callback2, 500);

  AssertJUnit.assertTrue(callback2.isTimedOut());

  cr.setPartition("TestDB_17");
  AsyncCallback callback3 = new MockAsyncCallback();
  int messageSent3 =
      _participants[0].getMessagingService().sendAndWait(cr, msg, callback3, 10000);
  AssertJUnit.assertTrue(callback3.getMessageReplied().size() == _replica - 1);

  cr.setPartition("TestDB_15");
  AsyncCallback callback4 = new MockAsyncCallback();
  int messageSent4 =
      _participants[0].getMessagingService().sendAndWait(cr, msg, callback4, 10000);
  AssertJUnit.assertTrue(callback4.getMessageReplied().size() == _replica);

  cr.setPartitionState("SLAVE");
  AsyncCallback callback5 = new MockAsyncCallback();
  int messageSent5 =
      _participants[0].getMessagingService().sendAndWait(cr, msg, callback5, 10000);
  AssertJUnit.assertTrue(callback5.getMessageReplied().size() == _replica - 1);

  cr.setDataSource(DataSource.IDEALSTATES);
  AsyncCallback callback6 = new MockAsyncCallback();
  int messageSent6 =
      _participants[0].getMessagingService().sendAndWait(cr, msg, callback6, 10000);
  AssertJUnit.assertTrue(callback6.getMessageReplied().size() == _replica - 1);
}
 

private static Message newMsg() {
  Message msg = new Message(DUMMY_MSG_TYPE, UUID.randomUUID().toString());
  msg.setTgtSessionId("*");
  msg.setTgtName("localhost_12918");
  return msg;
}
 

@Test()
public void TestControllerMessage() {
  for (int i = 0; i < NODE_NR; i++) {
    TestMessagingHandlerFactory factory = new TestMessagingHandlerFactory();
    _participants[i].getMessagingService()
        .registerMessageHandlerFactory(factory.getMessageTypes(), factory);

  }
  String msgId = new UUID(123, 456).toString();
  Message msg = new Message(MessageType.CONTROLLER_MSG, msgId);
  msg.setMsgId(msgId);
  msg.setSrcName(_hostSrc);
  msg.setTgtSessionId("*");
  msg.setMsgState(MessageState.NEW);
  String para = "Testing messaging para";
  msg.getRecord().setSimpleField("TestMessagingPara", para);

  Criteria cr = new Criteria();
  cr.setInstanceName("*");
  cr.setRecipientInstanceType(InstanceType.CONTROLLER);
  cr.setSessionSpecific(false);
  cr.setClusterName(CLUSTER_NAME);

  AsyncCallback callback1 = new MockAsyncCallback();
  int messagesSent =
      _adminController.getMessagingService().sendAndWait(cr, msg, callback1, 10000);

  AssertJUnit.assertTrue(callback1.getMessageReplied().get(0).getRecord()
      .getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ControllerResult")
      .contains(_hostSrc));
  AssertJUnit.assertEquals(callback1.getMessageReplied().size(), 1);

  msgId = UUID.randomUUID().toString();
  msg.setMsgId(msgId);
  cr.setPartition("TestDB_17");
  AsyncCallback callback2 = new MockAsyncCallback();
  messagesSent = _adminController.getMessagingService().sendAndWait(cr, msg, callback2, 10000);

  AssertJUnit.assertTrue(callback2.getMessageReplied().get(0).getRecord()
      .getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ControllerResult")
      .contains(_hostSrc));

  AssertJUnit.assertEquals(callback2.getMessageReplied().size(), 1);

  msgId = UUID.randomUUID().toString();
  msg.setMsgId(msgId);
  cr.setPartitionState("SLAVE");
  AsyncCallback callback3 = new MockAsyncCallback();
  messagesSent = _adminController.getMessagingService().sendAndWait(cr, msg, callback3, 10000);
  AssertJUnit.assertTrue(callback3.getMessageReplied().get(0).getRecord()
      .getMapField(Message.Attributes.MESSAGE_RESULT.toString()).get("ControllerResult")
      .contains(_hostSrc));

  AssertJUnit.assertEquals(callback3.getMessageReplied().size(), 1);
}
 

@Test
public void testNoWriteReadStateForRemovedMessage()
    throws NoSuchMethodException, InvocationTargetException, IllegalAccessException {
  HelixTaskExecutor executor = new HelixTaskExecutor();
  HelixManager manager = new MockClusterManager();
  TestMessageHandlerFactory factory = new TestMessageHandlerFactory();

  for (String type : factory.getMessageTypes()) {
    executor.registerMessageHandlerFactory(type, factory);
  }

  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  String instanceName = "someInstance";

  List<String> messageIds = new ArrayList<>();
  List<Message> messages = new ArrayList<>();
  int nMsgs1 = 5;
  for (int i = 0; i < nMsgs1; i++) {
    Message msg = new Message(factory.getMessageTypes().get(0), UUID.randomUUID().toString());
    msg.setTgtSessionId(manager.getSessionId());
    msg.setTgtName("Localhost_1123");
    msg.setSrcName("127.101.1.23_2234");
    msg.setCorrelationId(UUID.randomUUID().toString());
    accessor.setProperty(keyBuilder.message(instanceName, msg.getId()), msg);
    messageIds.add(msg.getId());
    messages.add(msg);
  }

  Method updateMessageState = HelixTaskExecutor.class
      .getDeclaredMethod("updateMessageState", List.class, HelixDataAccessor.class, String.class);
  updateMessageState.setAccessible(true);

  updateMessageState.invoke(executor, messages, accessor, instanceName);
  Assert.assertEquals(accessor.getChildNames(keyBuilder.messages(instanceName)).size(), nMsgs1);

  accessor.removeProperty(keyBuilder.message(instanceName, messageIds.get(0)));
  System.out.println(accessor.getChildNames(keyBuilder.messages(instanceName)).size());
  updateMessageState.invoke(executor, messages, accessor, instanceName);
  Assert
      .assertEquals(accessor.getChildNames(keyBuilder.messages(instanceName)).size(), nMsgs1 - 1);
}
 

@Test()
public void TestParticipantErrorMessageSend() {
  String participant1 = "localhost_" + START_PORT;
  String participant2 = "localhost_" + (START_PORT + 1);

  Message errorMessage1 =
      new Message(MessageType.PARTICIPANT_ERROR_REPORT, UUID.randomUUID().toString());
  errorMessage1.setTgtSessionId("*");
  errorMessage1.getRecord().setSimpleField(
      DefaultParticipantErrorMessageHandlerFactory.ACTIONKEY,
      ActionOnError.DISABLE_INSTANCE.toString());
  Criteria recipientCriteria = new Criteria();
  recipientCriteria.setRecipientInstanceType(InstanceType.CONTROLLER);
  recipientCriteria.setSessionSpecific(false);
  _participants[0].getMessagingService().send(recipientCriteria,
      errorMessage1);

  Message errorMessage2 =
      new Message(MessageType.PARTICIPANT_ERROR_REPORT, UUID.randomUUID().toString());
  errorMessage2.setTgtSessionId("*");
  errorMessage2.setResourceName("TestDB");
  errorMessage2.setPartitionName("TestDB_14");
  errorMessage2.getRecord().setSimpleField(
      DefaultParticipantErrorMessageHandlerFactory.ACTIONKEY,
      ActionOnError.DISABLE_PARTITION.toString());
  Criteria recipientCriteria2 = new Criteria();
  recipientCriteria2.setRecipientInstanceType(InstanceType.CONTROLLER);
  recipientCriteria2.setSessionSpecific(false);
  _participants[1].getMessagingService().send(recipientCriteria2,
      errorMessage2);

  try {
    Thread.sleep(1500);
  } catch (InterruptedException e) {
    LOG.error("Interrupted sleep", e);
  }

  boolean result =
      ClusterStateVerifier.verifyByZkCallback(new BestPossAndExtViewZkVerifier(ZK_ADDR,
          CLUSTER_NAME));
  Assert.assertTrue(result);
  Builder kb = _participants[1].getHelixDataAccessor().keyBuilder();
  ExternalView externalView =
      _participants[1].getHelixDataAccessor().getProperty(
          kb.externalView("TestDB"));

  for (String partitionName : externalView.getRecord().getMapFields().keySet()) {
    for (String hostName : externalView.getRecord().getMapField(partitionName).keySet()) {
      if (hostName.equals(participant1)) {
        Assert.assertTrue(externalView.getRecord().getMapField(partitionName).get(hostName)
            .equalsIgnoreCase("OFFLINE"));
      }
    }
  }
  Assert.assertTrue(externalView.getRecord().getMapField("TestDB_14").get(participant2)
      .equalsIgnoreCase("OFFLINE"));
}