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

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

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

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

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

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

}
 

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

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

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

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