Java Code Examples for org.apache.helix.Criteria#setRecipientInstanceType()

@VisibleForTesting
public static void sendUserDefinedMessage(String messageSubType, String messageVal, String messageId,
    InstanceType instanceType, HelixManager helixManager, Logger logger) {
  Criteria criteria = new Criteria();
  criteria.setInstanceName("%");
  criteria.setResource("%");
  criteria.setPartition("%");
  criteria.setPartitionState("%");
  criteria.setRecipientInstanceType(instanceType);
  criteria.setSessionSpecific(true);

  Message message = new Message(Message.MessageType.USER_DEFINE_MSG.toString(), messageId);
  message.setMsgSubType(messageSubType);
  message.setAttribute(Message.Attributes.INNER_MESSAGE, messageVal);
  message.setMsgState(Message.MessageState.NEW);
  message.setTgtSessionId("*");

  int messagesSent = helixManager.getMessagingService().send(criteria, message);
  if (messagesSent == 0) {
    logger.error(String.format("Failed to send the %s message to the participants", message));
  }
}
 

public int reloadAllSegments(String tableNameWithType) {
  LOGGER.info("Sending reload message for table: {}", tableNameWithType);

  Criteria recipientCriteria = new Criteria();
  recipientCriteria.setRecipientInstanceType(InstanceType.PARTICIPANT);
  recipientCriteria.setInstanceName("%");
  recipientCriteria.setResource(tableNameWithType);
  recipientCriteria.setSessionSpecific(true);
  SegmentReloadMessage segmentReloadMessage = new SegmentReloadMessage(tableNameWithType, null);
  ClusterMessagingService messagingService = _helixZkManager.getMessagingService();

  // Infinite timeout on the recipient
  int timeoutMs = -1;
  int numMessagesSent = messagingService.send(recipientCriteria, segmentReloadMessage, null, timeoutMs);
  if (numMessagesSent > 0) {
    LOGGER.info("Sent {} reload messages for table: {}", numMessagesSent, tableNameWithType);
  } else {
    LOGGER.warn("No reload message sent for table: {}", tableNameWithType);
  }

  return numMessagesSent;
}
 

public int reloadSegment(String tableNameWithType, String segmentName) {
  LOGGER.info("Sending reload message for segment: {} in table: {}", segmentName, tableNameWithType);

  Criteria recipientCriteria = new Criteria();
  recipientCriteria.setRecipientInstanceType(InstanceType.PARTICIPANT);
  recipientCriteria.setInstanceName("%");
  recipientCriteria.setResource(tableNameWithType);
  recipientCriteria.setPartition(segmentName);
  recipientCriteria.setSessionSpecific(true);
  SegmentReloadMessage segmentReloadMessage = new SegmentReloadMessage(tableNameWithType, segmentName);
  ClusterMessagingService messagingService = _helixZkManager.getMessagingService();

  // Infinite timeout on the recipient
  int timeoutMs = -1;
  int numMessagesSent = messagingService.send(recipientCriteria, segmentReloadMessage, null, timeoutMs);
  if (numMessagesSent > 0) {
    LOGGER.info("Sent {} reload messages for segment: {} in table: {}", numMessagesSent, segmentName,
        tableNameWithType);
  } else {
    LOGGER.warn("No reload message sent for segment: {} in table: {}", segmentName, tableNameWithType);
  }
  return numMessagesSent;
}
 

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

private void sendTableConfigRefreshMessage(String tableNameWithType) {
  TableConfigRefreshMessage tableConfigRefreshMessage = new TableConfigRefreshMessage(tableNameWithType);

  // Send table config refresh message to brokers
  Criteria recipientCriteria = new Criteria();
  recipientCriteria.setRecipientInstanceType(InstanceType.PARTICIPANT);
  recipientCriteria.setInstanceName("%");
  recipientCriteria.setResource(Helix.BROKER_RESOURCE_INSTANCE);
  recipientCriteria.setSessionSpecific(true);
  recipientCriteria.setPartition(tableNameWithType);
  // Send message with no callback and infinite timeout on the recipient
  int numMessagesSent =
      _helixZkManager.getMessagingService().send(recipientCriteria, tableConfigRefreshMessage, null, -1);
  if (numMessagesSent > 0) {
    // TODO: Would be nice if we can get the name of the instances to which messages were sent
    LOGGER.info("Sent {} table config refresh messages to brokers for table: {}", numMessagesSent, tableNameWithType);
  } else {
    LOGGER.warn("No table config refresh message sent to brokers for table: {}", tableNameWithType);
  }
}
 

@VisibleForTesting
void sendShutdownRequest() {
  final Criteria criteria = new Criteria();
  criteria.setInstanceName("%");
  criteria.setResource("%");
  criteria.setPartition("%");
  criteria.setPartitionState("%");
  criteria.setRecipientInstanceType(InstanceType.CONTROLLER);
  criteria.setSessionSpecific(true);

  final 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("*");

  // Wait for 5 minutes
  final int timeout = 300000;

  // Send shutdown request to Cluster master, which will send shutdown request to workers
  // Upon receiving shutdown response from workers, master will shut itself down and call back shutdownASG()
  final int messagesSent = this.helixManager.getMessagingService().send(criteria, shutdownRequest,
      shutdownASG(),timeout);
  if (messagesSent == 0) {
    LOGGER.error(String.format("Failed to send the %s message to the controller", shutdownRequest.getMsgSubType()));
  }
}
 

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

@Transition(from = "OFFLINE", to = "SLAVE")
public void offlineToSlave(Message message, NotificationContext context) {
  System.out.println("BootstrapProcess.BootstrapStateModel.offlineToSlave()");
  HelixManager manager = context.getManager();
  ClusterMessagingService messagingService = manager.getMessagingService();
  Message requestBackupUriRequest =
      new Message(MessageType.USER_DEFINE_MSG, UUID.randomUUID().toString());
  requestBackupUriRequest.setMsgSubType(BootstrapProcess.REQUEST_BOOTSTRAP_URL);
  requestBackupUriRequest.setMsgState(MessageState.NEW);
  Criteria recipientCriteria = new Criteria();
  recipientCriteria.setInstanceName("*");
  recipientCriteria.setRecipientInstanceType(InstanceType.PARTICIPANT);
  recipientCriteria.setResource(message.getResourceName());
  recipientCriteria.setPartition(message.getPartitionName());
  recipientCriteria.setSessionSpecific(true);
  // wait for 30 seconds
  int timeout = 30000;
  BootstrapReplyHandler responseHandler = new BootstrapReplyHandler();

  int sentMessageCount =
      messagingService.sendAndWait(recipientCriteria, requestBackupUriRequest, responseHandler,
          timeout);
  if (sentMessageCount == 0) {
    // could not find any other node hosting the partition
  } else if (responseHandler.getBootstrapUrl() != null) {
    System.out.println("Got bootstrap url:" + responseHandler.getBootstrapUrl());
    System.out.println("Got backup time:" + responseHandler.getBootstrapTime());
    // Got the url fetch it
  } else {
    // Either go to error state
    // throw new Exception("Cant find backup/bootstrap data");
    // Request some node to start backup process
  }
}
 

@VisibleForTesting
void sendShutdownRequest() {
  Criteria criteria = new Criteria();
  criteria.setInstanceName("%");
  criteria.setResource("%");
  criteria.setPartition("%");
  criteria.setPartitionState("%");
  criteria.setRecipientInstanceType(InstanceType.PARTICIPANT);
  // #HELIX-0.6.7-WORKAROUND
  // Add this back when messaging to instances is ported to 0.6 branch
  //criteria.setDataSource(Criteria.DataSource.LIVEINSTANCES);
  criteria.setSessionSpecific(true);

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

  // Wait for 5 minutes
  final int timeout = 300000;

  // #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, shutdownRequest,
  //    new NoopReplyHandler(), timeout);
  GobblinHelixMessagingService messagingService = new GobblinHelixMessagingService(this.multiManager.getJobClusterHelixManager());

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

/**
 * Attempt to send a message to refresh the new segment. We do not wait for any acknowledgements.
 * The message is sent as session-specific, so if a new zk session is created (e.g. server restarts)
 * it will not get the message.
 */
private void sendSegmentRefreshMessage(String offlineTableName, String segmentName) {
  SegmentRefreshMessage segmentRefreshMessage = new SegmentRefreshMessage(offlineTableName, segmentName);

  // Send segment refresh message to servers
  Criteria recipientCriteria = new Criteria();
  recipientCriteria.setRecipientInstanceType(InstanceType.PARTICIPANT);
  recipientCriteria.setInstanceName("%");
  recipientCriteria.setResource(offlineTableName);
  recipientCriteria.setPartition(segmentName);
  recipientCriteria.setSessionSpecific(true);
  ClusterMessagingService messagingService = _helixZkManager.getMessagingService();
  // Send message with no callback and infinite timeout on the recipient
  int numMessagesSent = messagingService.send(recipientCriteria, segmentRefreshMessage, null, -1);
  if (numMessagesSent > 0) {
    // TODO: Would be nice if we can get the name of the instances to which messages were sent
    LOGGER.info("Sent {} segment refresh messages to servers for segment: {} of table: {}", numMessagesSent,
        segmentName, offlineTableName);
  } else {
    // May be the case when none of the servers are up yet. That is OK, because when they come up they will get the
    // new version of the segment.
    LOGGER.warn("No segment refresh message sent to servers for segment: {} of table: {}", segmentName,
        offlineTableName);
  }

  // Send segment refresh message to brokers
  recipientCriteria.setResource(Helix.BROKER_RESOURCE_INSTANCE);
  recipientCriteria.setPartition(offlineTableName);
  numMessagesSent = messagingService.send(recipientCriteria, segmentRefreshMessage, null, -1);
  if (numMessagesSent > 0) {
    // TODO: Would be nice if we can get the name of the instances to which messages were sent
    LOGGER.info("Sent {} segment refresh messages to brokers for segment: {} of table: {}", numMessagesSent,
        segmentName, offlineTableName);
  } else {
    LOGGER.warn("No segment refresh message sent to brokers for segment: {} of table: {}", segmentName,
        offlineTableName);
  }
}
 

@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 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 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 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 TestMultiMessageCriteria() {
  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(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.setClusterName(CLUSTER_NAME);

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

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

  AsyncCallback callback2 = new MockAsyncCallback();
  int messageSent2 = _adminController.getMessagingService().sendAndWait(cr, msg, callback2, 500);

  AssertJUnit.assertTrue(callback2.isTimedOut());

  cr.setPartition("TestDB_17");
  AsyncCallback callback3 = new MockAsyncCallback();
  int messageSent3 =
      _adminController.getMessagingService().sendAndWait(cr, msg, callback3, 10000);
  AssertJUnit.assertEquals(_replica, callback3.getMessageReplied().size());

  cr.setPartition("TestDB_15");
  AsyncCallback callback4 = new MockAsyncCallback();
  int messageSent4 =
      _adminController.getMessagingService().sendAndWait(cr, msg, callback4, 10000);
  AssertJUnit.assertEquals(_replica, callback4.getMessageReplied().size());

  cr.setPartitionState("SLAVE");
  AsyncCallback callback5 = new MockAsyncCallback();
  int messageSent5 =
      _adminController.getMessagingService().sendAndWait(cr, msg, callback5, 10000);
  AssertJUnit.assertEquals(_replica - 1, callback5.getMessageReplied().size());

  cr.setDataSource(DataSource.IDEALSTATES);
  AsyncCallback callback6 = new MockAsyncCallback();
  int messageSent6 =
      _adminController.getMessagingService().sendAndWait(cr, msg, callback6, 10000);
  AssertJUnit.assertEquals(_replica - 1, callback6.getMessageReplied().size());
}
 

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

@Test(dependsOnMethods = "testConfigureResolver")
public void testExcessiveTopStateResolver() throws InterruptedException {
  BestPossibleExternalViewVerifier verifier =
      new BestPossibleExternalViewVerifier.Builder(CLUSTER_NAME).setZkClient(_gZkClient).build();
  Assert.assertTrue(verifier.verify());

  // 1. Find a partition with a MASTER replica and a SLAVE replica
  HelixAdmin admin = new ZKHelixAdmin.Builder().setZkAddress(ZK_ADDR).build();
  ExternalView ev = admin.getResourceExternalView(CLUSTER_NAME, TEST_DB);
  String targetPartition = ev.getPartitionSet().iterator().next();
  Map<String, String> partitionAssignment = ev.getStateMap(targetPartition);
  String slaveHost = partitionAssignment.entrySet().stream()
      .filter(entry -> entry.getValue().equals(MasterSlaveSMD.States.SLAVE.name())).findAny()
      .get().getKey();
  long previousMasterUpdateTime =
      getTopStateUpdateTime(ev, targetPartition, MasterSlaveSMD.States.MASTER.name());

  // Build SLAVE to MASTER message
  String msgId = new UUID(123, 456).toString();
  Message msg = createMessage(Message.MessageType.STATE_TRANSITION, msgId,
      MasterSlaveSMD.States.SLAVE.name(), MasterSlaveSMD.States.MASTER.name(), TEST_DB,
      slaveHost);
  msg.setStateModelDef(MasterSlaveSMD.name);

  Criteria cr = new Criteria();
  cr.setInstanceName(slaveHost);
  cr.setRecipientInstanceType(InstanceType.PARTICIPANT);
  cr.setSessionSpecific(true);
  cr.setPartition(targetPartition);
  cr.setResource(TEST_DB);
  cr.setClusterName(CLUSTER_NAME);

  AsyncCallback callback = new AsyncCallback() {
    @Override
    public void onTimeOut() {
      Assert.fail("The test state transition timeout.");
    }

    @Override
    public void onReplyMessage(Message message) {
      Assert.assertEquals(message.getMsgState(), Message.MessageState.READ);
    }
  };

  // 2. Send the SLAVE to MASTER message to the SLAVE host to make abnormal partition states.

  // 2.A. Without resolver, the fixing is not completely done by the default rebalancer logic.
  _controller.getMessagingService()
      .sendAndWait(cr, msg, callback, (int) TestHelper.WAIT_DURATION);
  Thread.sleep(DEFAULT_REBALANCE_PROCESSING_WAIT_TIME);
  // Wait until the partition status is fixed, verify if the result is as expected
  verifier =
      new BestPossibleExternalViewVerifier.Builder(CLUSTER_NAME).setZkClient(_gZkClient).build();
  Assert.assertTrue(verifier.verifyByPolling());
  ev = admin.getResourceExternalView(CLUSTER_NAME, TEST_DB);
  Assert.assertEquals(ev.getStateMap(targetPartition).values().stream()
      .filter(state -> state.equals(MasterSlaveSMD.States.MASTER.name())).count(), 1);
  // Since the resolver is not used in the auto default fix process, there is no update on the
  // original master. So if there is any data issue, it was not fixed.
  long currentMasterUpdateTime =
      getTopStateUpdateTime(ev, targetPartition, MasterSlaveSMD.States.MASTER.name());
  Assert.assertFalse(currentMasterUpdateTime > previousMasterUpdateTime);

  // 2.B. with resolver configured, the fixing is complete.
  ConfigAccessor configAccessor = new ConfigAccessor.Builder().setZkAddress(ZK_ADDR).build();
  ClusterConfig clusterConfig = configAccessor.getClusterConfig(CLUSTER_NAME);
  clusterConfig.setAbnormalStateResolverMap(
      ImmutableMap.of(MasterSlaveSMD.name, ExcessiveTopStateResolver.class.getName()));
  configAccessor.setClusterConfig(CLUSTER_NAME, clusterConfig);

  _controller.getMessagingService()
      .sendAndWait(cr, msg, callback, (int) TestHelper.WAIT_DURATION);
  Thread.sleep(DEFAULT_REBALANCE_PROCESSING_WAIT_TIME);
  // Wait until the partition status is fixed, verify if the result is as expected
  Assert.assertTrue(verifier.verifyByPolling());
  ev = admin.getResourceExternalView(CLUSTER_NAME, TEST_DB);
  Assert.assertEquals(ev.getStateMap(targetPartition).values().stream()
      .filter(state -> state.equals(MasterSlaveSMD.States.MASTER.name())).count(), 1);
  // Now the resolver is used in the auto fix process, the original master has also been refreshed.
  // The potential data issue has been fixed in this process.
  currentMasterUpdateTime =
      getTopStateUpdateTime(ev, targetPartition, MasterSlaveSMD.States.MASTER.name());
  Assert.assertTrue(currentMasterUpdateTime > previousMasterUpdateTime);

  // Reset the resolver map
  clusterConfig = configAccessor.getClusterConfig(CLUSTER_NAME);
  clusterConfig.setAbnormalStateResolverMap(Collections.emptyMap());
  configAccessor.setClusterConfig(CLUSTER_NAME, clusterConfig);
}
 

@Test()
public void TestMessageSend() {
  HelixManager manager = new MockHelixManager();
  DefaultMessagingService svc = new DefaultMessagingService(manager);
  TestMessageHandlerFactory factory = new TestMessageHandlerFactory();
  svc.registerMessageHandlerFactory(factory.getMessageType(), factory);

  Criteria recipientCriteria = new Criteria();
  recipientCriteria.setInstanceName("localhost_12919");
  recipientCriteria.setRecipientInstanceType(InstanceType.PARTICIPANT);
  recipientCriteria.setSelfExcluded(true);

  Message template = new Message(factory.getMessageType(), UUID.randomUUID().toString());
  AssertJUnit.assertEquals(0, svc.send(recipientCriteria, template));

  recipientCriteria.setSelfExcluded(false);
  AssertJUnit.assertEquals(1, svc.send(recipientCriteria, template));

  // all instances, all partitions
  recipientCriteria.setSelfExcluded(false);
  recipientCriteria.setInstanceName("%");
  recipientCriteria.setResource("DB");
  recipientCriteria.setPartition("%");
  AssertJUnit.assertEquals(200, svc.send(recipientCriteria, template));

  // all instances, all partitions, use * instead of %
  recipientCriteria.setSelfExcluded(false);
  recipientCriteria.setInstanceName("*");
  recipientCriteria.setResource("DB");
  recipientCriteria.setPartition("*");
  AssertJUnit.assertEquals(200, svc.send(recipientCriteria, template));

  // tail pattern
  recipientCriteria.setSelfExcluded(false);
  recipientCriteria.setInstanceName("localhost%");
  recipientCriteria.setResource("DB");
  recipientCriteria.setPartition("%");
  AssertJUnit.assertEquals(200, svc.send(recipientCriteria, template));

  // exclude this instance, send to all others for all partitions
  recipientCriteria.setSelfExcluded(true);
  recipientCriteria.setInstanceName("%");
  recipientCriteria.setResource("DB");
  recipientCriteria.setPartition("%");
  AssertJUnit.assertEquals(159, svc.send(recipientCriteria, template));

  // single instance, all partitions
  recipientCriteria.setSelfExcluded(true);
  recipientCriteria.setInstanceName("localhost_12920");
  recipientCriteria.setResource("DB");
  recipientCriteria.setPartition("%");
  AssertJUnit.assertEquals(39, svc.send(recipientCriteria, template));

  // single character wildcards
  recipientCriteria.setSelfExcluded(true);
  recipientCriteria.setInstanceName("l_calhost_12_20");
  recipientCriteria.setResource("DB");
  recipientCriteria.setPartition("%");
  AssertJUnit.assertEquals(39, svc.send(recipientCriteria, template));

  // head pattern
  recipientCriteria.setSelfExcluded(true);
  recipientCriteria.setInstanceName("%12920");
  recipientCriteria.setResource("DB");
  recipientCriteria.setPartition("%");
  AssertJUnit.assertEquals(39, svc.send(recipientCriteria, template));

  // middle pattern
  recipientCriteria.setSelfExcluded(true);
  recipientCriteria.setInstanceName("l%_12920");
  recipientCriteria.setResource("DB");
  recipientCriteria.setPartition("%");
  AssertJUnit.assertEquals(39, svc.send(recipientCriteria, template));

  // send to a controller
  recipientCriteria.setSelfExcluded(true);
  recipientCriteria.setInstanceName("localhost_12920");
  recipientCriteria.setRecipientInstanceType(InstanceType.CONTROLLER);
  recipientCriteria.setResource("DB");
  recipientCriteria.setPartition("%");
  AssertJUnit.assertEquals(1, svc.send(recipientCriteria, template));
}