org.apache.helix.Criteria Java Examples

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

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

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

@Override
// TODO if the manager is not Participant or Controller, no reply, so should fail immediately
public int sendAndWait(Criteria recipientCriteria, Message message, AsyncCallback asyncCallback,
    int timeOut, int retryCount) {
  int messagesSent = send(recipientCriteria, message, asyncCallback, timeOut, retryCount);
  if (messagesSent > 0) {
    synchronized (asyncCallback) {
      while (!asyncCallback.isDone() && !asyncCallback.isTimedOut()) {
        try {
          asyncCallback.wait();
        } catch (InterruptedException e) {
          _logger.error(e.toString());
          asyncCallback.setInterrupted(true);
          break;
        }
      }
    }
  } else {
    _logger.warn("No messages sent. For Criteria:" + recipientCriteria);
  }
  return messagesSent;
}
 

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

public Map<InstanceType, List<Message>> generateMessage(final Criteria recipientCriteria,
    final Message message) {
  Map<InstanceType, List<Message>> messagesToSendMap = new HashMap<InstanceType, List<Message>>();
  InstanceType instanceType = recipientCriteria.getRecipientInstanceType();

  HelixDataAccessor targetDataAccessor = getRecipientDataAccessor(recipientCriteria);

    List<Message> messages = Collections.EMPTY_LIST;
    if (instanceType == InstanceType.CONTROLLER) {
      messages = generateMessagesForController(message);
    } else if (instanceType == InstanceType.PARTICIPANT) {
      messages =
          generateMessagesForParticipant(recipientCriteria, message, targetDataAccessor);
    }
    messagesToSendMap.put(instanceType, messages);
    return messagesToSendMap;
}
 

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

/**
 * 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 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> generateMessagesForParticipant(Criteria recipientCriteria, Message message,
    HelixDataAccessor targetDataAccessor) {
  List<Message> messages = new ArrayList<Message>();
  List<Map<String, String>> matchedList =
      _evaluator.evaluateCriteria(recipientCriteria, targetDataAccessor);

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

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

private HelixDataAccessor getRecipientDataAccessor(final Criteria recipientCriteria) {
  HelixDataAccessor dataAccessor = _manager.getHelixDataAccessor();
  String clusterName = recipientCriteria.getClusterName();
  if (clusterName != null && !clusterName.equals(_manager.getClusterName())) {
    // for cross cluster message, create new DataAccessor for sending message.
    /*
      TODO On frequent cross clsuter messaging request, keeping construct data accessor may cause
      performance issue. We should consider adding cache in this service or HelixManager. --JJ
     */
    dataAccessor = new ZKHelixDataAccessor(clusterName, dataAccessor.getBaseDataAccessor());
  }
  return dataAccessor;
}
 

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

/**
 * Check if a given row matches the specified criteria
 * @param criteria the criteria
 * @param row row of currently persisted data
 * @return true if it matches, false otherwise
 */
private boolean rowMatches(Criteria criteria, ZNRecordRow row) {
  String instanceName = normalizePattern(criteria.getInstanceName());
  String resourceName = normalizePattern(criteria.getResource());
  String partitionName = normalizePattern(criteria.getPartition());
  String partitionState = normalizePattern(criteria.getPartitionState());
  return (stringMatches(instanceName, Strings.nullToEmpty(row.getMapSubKey()))
      || stringMatches(instanceName, Strings.nullToEmpty(row.getRecordId())))
      && stringMatches(resourceName, Strings.nullToEmpty(row.getRecordId()))
      && stringMatches(partitionName, Strings.nullToEmpty(row.getMapKey()))
      && stringMatches(partitionState, Strings.nullToEmpty(row.getMapValue()));
}
 

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

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

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

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

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

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

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

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

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

/**
 * Check if a given row matches the specified criteria
 * @param criteria the criteria
 * @param row row of currently persisted data
 * @return true if it matches, false otherwise
 */

private boolean rowMatches(Criteria criteria, ZNRecordRow row) {
  String instanceName = normalizePattern(criteria.getInstanceName());
  String resourceName = normalizePattern(criteria.getResource());
  String partitionName = normalizePattern(criteria.getPartition());
  String partitionState = normalizePattern(criteria.getPartitionState());
  return (stringMatches(instanceName, Strings.nullToEmpty(row.getMapSubKey())) ||
      stringMatches(instanceName, Strings.nullToEmpty(row.getRecordId())))
      && stringMatches(resourceName, Strings.nullToEmpty(row.getRecordId()))
      && stringMatches(partitionName, Strings.nullToEmpty(row.getMapKey()))
      && stringMatches(partitionState, Strings.nullToEmpty(row.getMapValue()));
}
 

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

@Override
public Map<InstanceType, List<Message>> generateMessage(Criteria recipientCriteria,
    Message messageTemplate) {
  // TODO Auto-generated method stub
  return null;
}
 

@Override
public int sendAndWait(Criteria receipientCriteria, Message message,
    AsyncCallback callbackOnReply, int timeOut, int retryCount) {
  // TODO Auto-generated method stub
  return 0;
}
 

@Override
public int send(Criteria receipientCriteria, Message message, AsyncCallback callbackOnReply,
    int timeOut, int retryCount) {
  // TODO Auto-generated method stub
  return 0;
}
 

@Override
public int sendAndWait(Criteria receipientCriteria, Message message,
    AsyncCallback callbackOnReply, int timeOut) {
  // TODO Auto-generated method stub
  return 0;
}
 

@Override
public int send(Criteria receipientCriteria, Message message, AsyncCallback callbackOnReply,
    int timeOut) {
  // TODO Auto-generated method stub
  return 0;
}
 

@Override
public int send(Criteria recipientCriteria, Message message) {
  // TODO Auto-generated method stub
  return 0;
}
 

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