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

private void sendNopMessage() {
  if (_manager.isConnected()) {
    try {
      Message nopMsg = new Message(MessageType.NO_OP, UUID.randomUUID().toString());
      nopMsg.setSrcName(_manager.getInstanceName());

      HelixDataAccessor accessor = _manager.getHelixDataAccessor();
      Builder keyBuilder = accessor.keyBuilder();

      if (_manager.getInstanceType() == InstanceType.CONTROLLER
          || _manager.getInstanceType() == InstanceType.CONTROLLER_PARTICIPANT) {
        nopMsg.setTgtName(InstanceType.CONTROLLER.name());
        accessor.setProperty(keyBuilder.controllerMessage(nopMsg.getId()), nopMsg);
      }

      if (_manager.getInstanceType() == InstanceType.PARTICIPANT
          || _manager.getInstanceType() == InstanceType.CONTROLLER_PARTICIPANT) {
        nopMsg.setTgtName(_manager.getInstanceName());
        accessor.setProperty(keyBuilder.message(nopMsg.getTgtName(), nopMsg.getId()), nopMsg);
      }
      logger.info("Send NO_OP message to " + nopMsg.getTgtName() + ", msgId: " + nopMsg.getId());
    } catch (Exception e) {
      logger.error(e.toString());
    }
  }
}
 

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 void sendNopMessageInternal() {
  try {
    Message nopMsg = new Message(MessageType.NO_OP, UUID.randomUUID().toString());
    nopMsg.setSrcName(_manager.getInstanceName());

    HelixDataAccessor accessor = _manager.getHelixDataAccessor();
    Builder keyBuilder = accessor.keyBuilder();

    if (_manager.getInstanceType() == InstanceType.CONTROLLER
        || _manager.getInstanceType() == InstanceType.CONTROLLER_PARTICIPANT) {
      nopMsg.setTgtName(InstanceType.CONTROLLER.name());
      accessor.setProperty(keyBuilder.controllerMessage(nopMsg.getId()), nopMsg);
    }

    if (_manager.getInstanceType() == InstanceType.PARTICIPANT
        || _manager.getInstanceType() == InstanceType.CONTROLLER_PARTICIPANT) {
      nopMsg.setTgtName(_manager.getInstanceName());
      accessor.setProperty(keyBuilder.message(nopMsg.getTgtName(), nopMsg.getId()), nopMsg);
    }
  } catch (Exception e) {
    _logger.error(e.toString());
  }
}
 

@Test()
public void testRetryOnce() throws InterruptedException {
  // Logger.getRootLogger().setLevel(Level.INFO);

  // String p = "test_";
  // System.out.println(p.substring(p.lastIndexOf('_')+1));
  HelixTaskExecutor executor = new HelixTaskExecutor();
  HelixManager manager = new MockClusterManager();

  CancellableHandlerFactory factory = new CancellableHandlerFactory();
  for (String type : factory.getMessageTypes()) {
    executor.registerMessageHandlerFactory(type, factory);
  }
  NotificationContext changeContext = new NotificationContext(manager);

  List<Message> msgList = new ArrayList<Message>();

  // factory.reset();
  // msgList.clear();
  // Test the case that the message are executed for the second time
  int nMsgs2 = 4;
  for (int i = 0; i < nMsgs2; i++) {
    Message msg = new Message(factory.getMessageTypes().get(0), UUID.randomUUID().toString());
    msg.setTgtSessionId("*");
    msg.setTgtName("Localhost_1123");
    msg.setSrcName("127.101.1.23_2234");
    msg.setExecutionTimeout((i + 1) * 600);
    msg.setRetryCount(1);
    msgList.add(msg);
  }
  changeContext.setChangeType(HelixConstants.ChangeType.MESSAGE);
  executor.onMessage("someInstance", msgList, changeContext);
  Thread.sleep(3500);
  AssertJUnit.assertEquals(factory._processedMsgIds.size(), 3);
  AssertJUnit.assertTrue(msgList.get(0).getRecord().getSimpleField("Cancelcount").equals("2"));
  AssertJUnit.assertTrue(msgList.get(1).getRecord().getSimpleField("Cancelcount").equals("1"));
  AssertJUnit.assertEquals(factory._timedOutMsgIds.size(), 2);
  AssertJUnit.assertTrue(executor._taskMap.size() == 0);

}
 

@Test()
public void TestBlockingSendReceive() {
  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);
  cr.setClusterName(CLUSTER_NAME);

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

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

  AsyncCallback asyncCallback2 = new MockAsyncCallback();
  messagesSent = _adminController.getMessagingService().sendAndWait(cr, msg, asyncCallback2, 500);
  AssertJUnit.assertTrue(asyncCallback2.isTimedOut());
}
 

@Test()
public void testNoRetry() throws InterruptedException {
  // String p = "test_";
  // System.out.println(p.substring(p.lastIndexOf('_')+1));
  HelixTaskExecutor executor = new HelixTaskExecutor();
  HelixManager manager = new MockClusterManager();

  CancellableHandlerFactory factory = new CancellableHandlerFactory();
  for (String type : factory.getMessageTypes()) {
    executor.registerMessageHandlerFactory(type, factory);
  }
  NotificationContext changeContext = new NotificationContext(manager);

  List<Message> msgList = new ArrayList<Message>();
  int nMsgs2 = 4;
  // Test the case in which retry = 0
  for (int i = 0; i < nMsgs2; i++) {
    Message msg = new Message(factory.getMessageTypes().get(0), UUID.randomUUID().toString());
    msg.setTgtSessionId("*");
    msg.setTgtName("Localhost_1123");
    msg.setSrcName("127.101.1.23_2234");
    msg.setExecutionTimeout((i + 1) * 600);
    msgList.add(msg);
  }
  changeContext.setChangeType(HelixConstants.ChangeType.MESSAGE);
  executor.onMessage("someInstance", msgList, changeContext);

  Thread.sleep(4000);

  AssertJUnit.assertTrue(factory._handlersCreated == nMsgs2);
  AssertJUnit.assertEquals(factory._timedOutMsgIds.size(), 2);
  // AssertJUnit.assertFalse(msgList.get(0).getRecord().getSimpleFields().containsKey("TimeOut"));
  for (int i = 0; i < nMsgs2 - 2; i++) {
    if (factory.getMessageTypes().contains(msgList.get(i).getMsgType())) {
      AssertJUnit.assertTrue(msgList.get(i).getRecord().getSimpleFields()
          .containsKey("Cancelcount"));
      AssertJUnit.assertTrue(factory._timedOutMsgIds.containsKey(msgList.get(i).getId()));
    }
  }
}
 

@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());

}
 

private List<Message> generateMessagesForController(Message message) {
  List<Message> messages = new ArrayList<Message>();
  String id = (message.getMsgId() == null) ? UUID.randomUUID().toString() : message.getMsgId();
  Message newMessage = new Message(message.getRecord(), id);
  newMessage.setMsgId(id);
  newMessage.setSrcName(_manager.getInstanceName());
  newMessage.setTgtName(InstanceType.CONTROLLER.name());
  messages.add(newMessage);
  return messages;
}
 

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> generateMessagesForController(Message message) {
  List<Message> messages = new ArrayList<Message>();
  String id = UUID.randomUUID().toString();
  Message newMessage = new Message(message.getRecord(), id);
  newMessage.setMsgId(id);
  newMessage.setSrcName(_manager.getInstanceName());
  newMessage.setTgtName("Controller");
  messages.add(newMessage);
  return messages;
}
 

@Test
public void testStateTransitionCancellationMsg() throws InterruptedException {
  HelixTaskExecutor executor = new HelixTaskExecutor();
  HelixManager manager = new MockClusterManager();

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


  NotificationContext changeContext = new NotificationContext(manager);

  List<Message> msgList = new ArrayList<Message>();
  Message msg1 = new Message(Message.MessageType.STATE_TRANSITION, UUID.randomUUID().toString());
  msg1.setTgtSessionId("*");
  msg1.setPartitionName("P1");
  msg1.setResourceName("R1");
  msg1.setTgtName("Localhost_1123");
  msg1.setSrcName("127.101.1.23_2234");
  msg1.setFromState("SLAVE");
  msg1.setToState("MASTER");
  msgList.add(msg1);

  Message msg2 = new Message(Message.MessageType.STATE_TRANSITION_CANCELLATION, UUID.randomUUID().toString());
  msg2.setTgtSessionId("*");
  msg2.setPartitionName("P1");
  msg2.setResourceName("R1");
  msg2.setTgtName("Localhost_1123");
  msg2.setSrcName("127.101.1.23_2234");
  msg2.setFromState("SLAVE");
  msg2.setToState("MASTER");
  msgList.add(msg2);

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

  Thread.sleep(3000);
  AssertJUnit.assertEquals(cancelFactory._processedMsgIds.size(), 0);
  AssertJUnit.assertEquals(stateTransitionFactory._processedMsgIds.size(), 0);
}
 

@Test
public void testStateMismatch() throws InterruptedException {
  // String controllerName = CONTROLLER_PREFIX + "_0";

  HelixManager manager = _controller; // _startCMResultMap.get(controllerName)._manager;
  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  Builder kb = accessor.keyBuilder();
  ExternalView ev = accessor.getProperty(kb.externalView(TEST_DB));
  Map<String, LiveInstance> liveinstanceMap =
      accessor.getChildValuesMap(accessor.keyBuilder().liveInstances(), true);

  for (String instanceName : liveinstanceMap.keySet()) {
    String sessionid = liveinstanceMap.get(instanceName).getEphemeralOwner();
    for (String partition : ev.getPartitionSet()) {
      if (ev.getStateMap(partition).containsKey(instanceName)) {
        String uuid = UUID.randomUUID().toString();
        Message message = new Message(MessageType.STATE_TRANSITION, uuid);
        boolean rand = new Random().nextInt(10) > 5;
        if (ev.getStateMap(partition).get(instanceName).equals("MASTER")) {
          message.setSrcName(manager.getInstanceName());
          message.setTgtName(instanceName);
          message.setMsgState(MessageState.NEW);
          message.setPartitionName(partition);
          message.setResourceName(TEST_DB);
          message.setFromState(rand ? "SLAVE" : "OFFLINE");
          message.setToState(rand ? "MASTER" : "SLAVE");
          message.setTgtSessionId(sessionid);
          message.setSrcSessionId(manager.getSessionId());
          message.setStateModelDef("MasterSlave");
          message.setStateModelFactoryName("DEFAULT");
        } else if (ev.getStateMap(partition).get(instanceName).equals("SLAVE")) {
          message.setSrcName(manager.getInstanceName());
          message.setTgtName(instanceName);
          message.setMsgState(MessageState.NEW);
          message.setPartitionName(partition);
          message.setResourceName(TEST_DB);
          message.setFromState(rand ? "MASTER" : "OFFLINE");
          message.setToState(rand ? "SLAVE" : "SLAVE");
          message.setTgtSessionId(sessionid);
          message.setSrcSessionId(manager.getSessionId());
          message.setStateModelDef("MasterSlave");
          message.setStateModelFactoryName("DEFAULT");
        }
        accessor.setProperty(accessor.keyBuilder().message(instanceName, message.getMsgId()),
            message);
      }
    }
  }
  Thread.sleep(3000);
  ExternalView ev2 = accessor.getProperty(kb.externalView(TEST_DB));
  Assert.assertTrue(ev.equals(ev2));
}
 

@Test()
public void testCreateHandlerException() throws InterruptedException {
  System.out.println("START TestCMTaskExecutor.testCreateHandlerException()");
  HelixTaskExecutor executor = new HelixTaskExecutor();
  HelixManager manager = new MockClusterManager();

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

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

  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());
    msgList.add(msg);
  }
  Message exceptionMsg = new Message(factory.getMessageTypes().get(0), UUID.randomUUID().toString());
  exceptionMsg.setTgtSessionId(manager.getSessionId());
  exceptionMsg.setMsgSubType("EXCEPTION");
  exceptionMsg.setTgtName("Localhost_1123");
  exceptionMsg.setSrcName("127.101.1.23_2234");
  exceptionMsg.setCorrelationId(UUID.randomUUID().toString());
  msgList.add(exceptionMsg);

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

  Thread.sleep(1000);

  AssertJUnit.assertTrue(factory._processedMsgIds.size() == nMsgs1);
  AssertJUnit.assertTrue(factory._handlersCreated == nMsgs1);

  AssertJUnit.assertTrue(exceptionMsg.getMsgState() == MessageState.UNPROCESSABLE);
  System.out.println("END TestCMTaskExecutor.testCreateHandlerException()");
}
 

@Test()
public void testCMTaskExecutor() throws Exception {
  System.out.println("START TestCMTaskExecutor");
  String msgId = "TestMessageId";
  Message message = new Message(MessageType.TASK_REPLY, msgId);

  message.setMsgId(msgId);
  message.setSrcName("cm-instance-0");
  message.setTgtName("cm-instance-1");
  message.setTgtSessionId("1234");
  message.setFromState("Offline");
  message.setToState("Slave");
  message.setPartitionName("TestDB_0");
  message.setResourceName("TestDB");
  message.setStateModelDef("MasterSlave");

  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);

  MockHelixTaskExecutor executor = new MockHelixTaskExecutor();
  MockMasterSlaveStateModel stateModel = new MockMasterSlaveStateModel();
  executor.registerMessageHandlerFactory(MessageType.TASK_REPLY.name(),
      new AsyncCallbackService());

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

  StateModelFactory<MockMasterSlaveStateModel> stateModelFactory = new StateModelFactory<MockMasterSlaveStateModel>() {

    @Override
    public MockMasterSlaveStateModel createNewStateModel(String resource, String partitionName) {
      // TODO Auto-generated method stub
      return new MockMasterSlaveStateModel();
    }

  };
  HelixStateTransitionHandler handler =
      new HelixStateTransitionHandler(stateModelFactory, stateModel, message, context,
          currentStateDelta);

  HelixTask task = new HelixTask(message, context, handler, executor);
  executor.scheduleTask(task);
  for (int i = 0; i < 10; i++) {
    if (!executor.isDone(task.getTaskId())) {
      Thread.sleep(500);
    }
  }
  AssertJUnit.assertTrue(stateModel.stateModelInvoked);
  System.out.println("END TestCMTaskExecutor");
}
 

@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);
}
 

@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 testMessageReadOptimization() throws InterruptedException {
  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();

  List<String> messageIds = 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("someInstance", msg.getId()), msg);
    messageIds.add(msg.getId());
  }

  NotificationContext changeContext = new NotificationContext(manager);
  changeContext.setChangeType(HelixConstants.ChangeType.MESSAGE);

  // Simulate read message already, then processing message. Should read and handle no message.
  executor._knownMessageIds.addAll(messageIds);
  executor.onMessage("someInstance", Collections.EMPTY_LIST, changeContext);
  Thread.sleep(3000);
  AssertJUnit.assertEquals(0, factory._processedMsgIds.size());
  executor._knownMessageIds.clear();

  // Processing message normally
  executor.onMessage("someInstance", Collections.EMPTY_LIST, changeContext);
  Thread.sleep(3000);
  AssertJUnit.assertEquals(nMsgs1, factory._processedMsgIds.size());
  // After all messages are processed, _knownMessageIds should be empty.
  Assert.assertTrue(executor._knownMessageIds.isEmpty());
}
 

@Test()
public void testSchedulerMsgUsingQueue() throws Exception {
  // Logger.getRootLogger().setLevel(Level.INFO);
  _factory._results.clear();
  Thread.sleep(2000);
  HelixManager manager = null;
  for (int i = 0; i < NODE_NR; i++) {
    _participants[i].getMessagingService().registerMessageHandlerFactory(
        _factory.getMessageType(), _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");
  schedulerMessage.getRecord().setSimpleField(
      DefaultSchedulerMessageHandlerFactory.SCHEDULER_TASK_QUEUE, "TestSchedulerMsg");
  // Template for the individual message sent to each participant
  Message msg = new Message(_factory.getMessageType(), "Template");
  msg.setTgtSessionId("*");
  msg.setMsgState(MessageState.NEW);

  // Criteria to send individual messages
  Criteria cr = new Criteria();
  cr.setInstanceName("localhost_%");
  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();
  PropertyKey.Builder keyBuilder = helixDataAccessor.keyBuilder();
  helixDataAccessor.createControllerMessage(schedulerMessage);

  for (int i = 0; i < 30; i++) {
    Thread.sleep(2000);
    if (_PARTITIONS == _factory._results.size()) {
      break;
    }
  }

  Assert.assertEquals(_PARTITIONS, _factory._results.size());
  PropertyKey controllerTaskStatus =
      keyBuilder.controllerTaskStatus(MessageType.SCHEDULER_MSG.name(),
          schedulerMessage.getMsgId());

  int messageResultCount = 0;
  for (int i = 0; i < 10; i++) {
    ZNRecord statusUpdate = helixDataAccessor.getProperty(controllerTaskStatus).getRecord();
    Assert.assertTrue(statusUpdate.getMapField("SentMessageCount").get("MessageCount")
        .equals("" + (_PARTITIONS * 3)));
    for (String key : statusUpdate.getMapFields().keySet()) {
      if (key.startsWith("MessageResult ")) {
        messageResultCount++;
      }
    }
    if (messageResultCount == _PARTITIONS * 3) {
      break;
    } else {
      Thread.sleep(2000);
    }
  }
  Assert.assertEquals(messageResultCount, _PARTITIONS * 3);
  int count = 0;
  for (Set<String> val : _factory._results.values()) {
    count += val.size();
  }
  Assert.assertEquals(count, _PARTITIONS * 3);

}
 

@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);
}