Java Code Examples for org.apache.helix.HelixManager#getHelixDataAccessor()

public ParticipantManager(HelixManager manager, RealmAwareZkClient zkclient, int sessionTimeout,
    LiveInstanceInfoProvider liveInstanceInfoProvider,
    List<PreConnectCallback> preConnectCallbacks, final String sessionId,
    HelixManagerProperty helixManagerProperty) {
  _zkclient = zkclient;
  _manager = manager;
  _clusterName = manager.getClusterName();
  _instanceName = manager.getInstanceName();
  _keyBuilder = new PropertyKey.Builder(_clusterName);
  _sessionId = sessionId;
  _sessionTimeout = sessionTimeout;
  _configAccessor = manager.getConfigAccessor();
  _instanceType = manager.getInstanceType();
  _helixAdmin = manager.getClusterManagmentTool();
  _dataAccessor = (ZKHelixDataAccessor) manager.getHelixDataAccessor();
  _messagingService = (DefaultMessagingService) manager.getMessagingService();
  _stateMachineEngine = manager.getStateMachineEngine();
  _liveInstanceInfoProvider = liveInstanceInfoProvider;
  _preConnectCallbacks = preConnectCallbacks;
  _helixManagerProperty = helixManagerProperty;
}
 

private void acquireLeadership(final HelixManager manager,
    ControllerManagerHelper controllerHelper) {
  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  PropertyKey leaderNodePropertyKey = accessor.keyBuilder().controllerLeader();

  LOG.info(manager.getInstanceName() + " tries to acquire leadership for cluster: " + manager
      .getClusterName());
  // Try to acquire leader and init the manager in any case.
  // Even when a leader node already exists, the election process shall still try to init the manager
  // in case it is the current leader.
  do {
    // Due to the possible carried over ZK events from the previous ZK session, the following
    // initialization might be triggered multiple times. So the operation must be idempotent.
    long start = System.currentTimeMillis();
    if (tryCreateController(manager)) {
      manager.getHelixDataAccessor().getBaseDataAccessor().reset();
      controllerHelper.addListenersToController(_controller);
      controllerHelper.startControllerTimerTasks();
      LOG.info("{} with session {} acquired leadership for cluster: {}, took: {}ms",
          manager.getInstanceName(), manager.getSessionId(), manager.getClusterName(),
          System.currentTimeMillis() - start);
    }
  } while (accessor.getProperty(leaderNodePropertyKey) == null);
}
 

private void updateHistory(HelixManager manager) {
  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  Builder keyBuilder = accessor.keyBuilder();
  final String clusterName = manager.getClusterName();
  final String instanceName = manager.getInstanceName();
  final String version = manager.getVersion();

  // Record a MaintenanceSignal history
  if (!accessor.getBaseDataAccessor().update(keyBuilder.controllerLeaderHistory().getPath(),
      oldRecord -> {
        if (oldRecord == null) {
          oldRecord = new ZNRecord(PropertyType.HISTORY.toString());
        }
        return new ControllerHistory(oldRecord).updateHistory(clusterName, instanceName,
            version);
      }, AccessOption.PERSISTENT)) {
    LOG.error("Failed to persist leader history to ZK!");
  }
}
 

public CallbackHandler(HelixManager manager, RealmAwareZkClient client, PropertyKey propertyKey,
    Object listener, EventType[] eventTypes, ChangeType changeType,
    HelixCallbackMonitor monitor) {
  if (listener == null) {
    throw new HelixException("listener could not be null");
  }

  if (monitor != null && !monitor.getChangeType().equals(changeType)) {
    throw new HelixException("The specified callback monitor is for different change type: "
        + monitor.getChangeType().name());
  }

  _manager = manager;
  _accessor = manager.getHelixDataAccessor();
  _zkClient = client;
  _propertyKey = propertyKey;
  _path = propertyKey.getPath();
  _listener = listener;
  _eventTypes = new HashSet<>(Arrays.asList(eventTypes));
  _changeType = changeType;
  _lastNotificationTimeStamp = new AtomicLong(System.nanoTime());
  _monitor = monitor;

  if (_changeType == MESSAGE || _changeType == MESSAGES_CONTROLLER || _changeType == CONTROLLER) {
    _watchChild = false;
  } else {
    _watchChild = true;
  }

  parseListenerProperties();

  init();
}
 

/**
 * @return true if the current manager created a new controller node successfully.
 */
private boolean tryCreateController(HelixManager manager) {
  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  Builder keyBuilder = accessor.keyBuilder();

  LiveInstance newLeader = new LiveInstance(manager.getInstanceName());
  newLeader.setLiveInstance(ManagementFactory.getRuntimeMXBean().getName());
  newLeader.setSessionId(manager.getSessionId());
  newLeader.setHelixVersion(manager.getVersion());
  try {
    if (accessor.createControllerLeader(newLeader)) {
      updateHistory(manager);
      return true;
    } else {
      LOG.info(
          "Unable to become leader probably because some other controller became the leader.");
    }
  } catch (Exception e) {
    LOG.error(
        "Exception when trying to updating leader record in cluster:" + manager.getClusterName()
            + ". Need to check again whether leader node has been created or not.", e);
  }

  LiveInstance currentLeader = accessor.getProperty(keyBuilder.controllerLeader());
  if (currentLeader != null) {
    String currentSession = currentLeader.getEphemeralOwner();
    LOG.info("Leader exists for cluster: " + manager.getClusterName() + ", currentLeader: "
        + currentLeader.getInstanceName() + ", leaderSessionId: " + currentSession);
    if (currentSession != null && currentSession.equals(newLeader.getEphemeralOwner())) {
      return true;
    } else {
      LOG.warn("The existing leader has a different session. Expected session Id: " + newLeader
          .getEphemeralOwner());
    }
  }
  return false;
}
 

private void setLeader(HelixManager manager) throws Exception {
  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  final LiveInstance leader = new LiveInstance(manager.getInstanceName());
  leader.setLiveInstance(ManagementFactory.getRuntimeMXBean().getName());
  leader.setSessionId(manager.getSessionId());
  leader.setHelixVersion(manager.getVersion());

  // Delete the current controller leader node so it will trigger leader election
  while (!manager.isLeader()) {
    accessor.getBaseDataAccessor()
        .remove(PropertyPathBuilder.controllerLeader(manager.getClusterName()),
            AccessOption.EPHEMERAL);
    Thread.sleep(50);
  }
}
 

/**
 * trigger dummy message handler and verify it's invoked
 * @param manager
 * @param handledMsgSet
 * @throws Exception
 */
private static void checkDummyMsgHandler(HelixManager manager, final Set<String> handledMsgSet)
    throws Exception {

  final Message aMsg = newMsg();
  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  accessor.setProperty(keyBuilder.message(manager.getInstanceName(), aMsg.getId()), aMsg);
  boolean result = TestHelper.verify(() -> handledMsgSet.contains(aMsg.getId()), 5 * 1000);
  Assert.assertTrue(result);
}
 

public void postHandleMessage() {
  if (_message.getBatchMessageMode() == true && _batchMsgWrapper != null) {
    _batchMsgWrapper.end(_message, _notificationContext);
  }

  // update currentState
  HelixManager manager = _notificationContext.getManager();
  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  ConcurrentHashMap<String, CurrentStateUpdate> csUpdateMap =
      (ConcurrentHashMap<String, CurrentStateUpdate>) _notificationContext
          .get(MapKey.CURRENT_STATE_UPDATE.toString());

  if (csUpdateMap != null) {
    Map<PropertyKey, CurrentState> csUpdate = mergeCurStateUpdate(csUpdateMap);

    // TODO: change to use asyncSet
    for (PropertyKey key : csUpdate.keySet()) {
      // logger.info("updateCS: " + key);
      // System.out.println("\tupdateCS: " + key.getPath() + ", " +
      // curStateMap.get(key));
      if(!accessor.updateProperty(key, csUpdate.get(key))) {
        LOG.error(
            "Fails to persist current state to ZK for key " + key);
      }
    }
  }
}
 

private static void printStatus(final HelixManager manager) {
  System.out.println("CLUSTER STATUS");
  HelixDataAccessor helixDataAccessor = manager.getHelixDataAccessor();
  Builder keyBuilder = helixDataAccessor.keyBuilder();
  System.out.println("External View \n"
      + helixDataAccessor.getProperty(keyBuilder.externalView("repository")));
}
 

@Test
public void testDataChange() throws Exception {
  // Create an admin and add LiveInstanceChange listener to it
  HelixManager adminManager =
      HelixManagerFactory.getZKHelixManager(clusterName, null, InstanceType.ADMINISTRATOR,
          ZK_ADDR);
  adminManager.connect();
  final BlockingQueue<List<LiveInstance>> changeList =
      new LinkedBlockingQueue<List<LiveInstance>>();

  adminManager.addLiveInstanceChangeListener(new LiveInstanceChangeListener() {
    @Override
    public void onLiveInstanceChange(List<LiveInstance> liveInstances,
        NotificationContext changeContext) {
      // The queue is basically unbounded, so shouldn't throw exception when calling
      // "add".
      changeList.add(deepCopy(liveInstances));
    }
  });

  // Check the initial condition
  List<LiveInstance> instances = changeList.poll(1, TimeUnit.SECONDS);
  Assert.assertNotNull(instances, "Expecting a list of live instance");
  Assert.assertTrue(instances.isEmpty(), "Expecting an empty list of live instance");
  // Join as participant, should trigger a live instance change event
  HelixManager manager =
      HelixManagerFactory.getZKHelixManager(clusterName, "localhost_54321",
          InstanceType.PARTICIPANT, ZK_ADDR);
  manager.connect();
  instances = changeList.poll(1, TimeUnit.SECONDS);
  Assert.assertNotNull(instances, "Expecting a list of live instance");
  Assert.assertEquals(instances.size(), 1, "Expecting one live instance");
  Assert.assertEquals(instances.get(0).getInstanceName(), manager.getInstanceName());
  // Update data in the live instance node, should trigger another live instance change
  // event
  HelixDataAccessor helixDataAccessor = manager.getHelixDataAccessor();
  PropertyKey propertyKey =
      helixDataAccessor.keyBuilder().liveInstance(manager.getInstanceName());
  LiveInstance instance = helixDataAccessor.getProperty(propertyKey);

  Map<String, String> map = new TreeMap<String, String>();
  map.put("k1", "v1");
  instance.getRecord().setMapField("test", map);
  Assert.assertTrue(helixDataAccessor.updateProperty(propertyKey, instance),
      "Failed to update live instance node");

  instances = changeList.poll(1, TimeUnit.SECONDS);
  Assert.assertNotNull(instances, "Expecting a list of live instance");
  Assert.assertEquals(instances.get(0).getRecord().getMapField("test"), map, "Wrong map data.");
  manager.disconnect();
  Thread.sleep(1000); // wait for callback finish

  instances = changeList.poll(1, TimeUnit.SECONDS);
  Assert.assertNotNull(instances, "Expecting a list of live instance");
  Assert.assertTrue(instances.isEmpty(), "Expecting an empty list of live instance");

  adminManager.disconnect();

}
 

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

}
 

public static IdealState getTableIdealState(HelixManager manager, String resourceName) {
  final HelixDataAccessor accessor = manager.getHelixDataAccessor();
  final Builder builder = accessor.keyBuilder();
  return accessor.getProperty(builder.idealStates(resourceName));
}
 

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

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

  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();

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

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

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

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

@Test
public void testInstanceAutoJoin() throws Exception {
  HelixManager manager = _participants[0];
  HelixDataAccessor accessor = manager.getHelixDataAccessor();

  _gSetupTool.addResourceToCluster(CLUSTER_NAME, db2, 60, "OnlineOffline",
      RebalanceMode.FULL_AUTO + "");

  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, db2, 1);
  String instance2 = "localhost_279699";
  // StartCMResult result = TestHelper.startDummyProcess(ZK_ADDR, CLUSTER_NAME, instance2);
  MockParticipantManager newParticipant =
      new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, instance2);
  newParticipant.syncStart();

  Thread.sleep(500);
  // Assert.assertFalse(result._thread.isAlive());
  Assert.assertTrue(null == manager.getHelixDataAccessor()
      .getProperty(accessor.keyBuilder().liveInstance(instance2)));

  ConfigScope scope = new ConfigScopeBuilder().forCluster(CLUSTER_NAME).build();

  manager.getConfigAccessor().set(scope, ZKHelixManager.ALLOW_PARTICIPANT_AUTO_JOIN, "true");

  newParticipant = new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, instance2);
  newParticipant.syncStart();

  Thread.sleep(500);
  // Assert.assertTrue(result._thread.isAlive() || result2._thread.isAlive());
  for (int i = 0; i < 20; i++) {
    if (null == manager.getHelixDataAccessor()
        .getProperty(accessor.keyBuilder().liveInstance(instance2))) {
      Thread.sleep(100);
    } else
      break;
  }
  Assert.assertTrue(null != manager.getHelixDataAccessor()
      .getProperty(accessor.keyBuilder().liveInstance(instance2)));

  newParticipant.syncStop();
}
 

protected void processEvent(ClusterEvent event, MessageOutput messageOutput) throws Exception {
  _eventId = event.getEventId();
  HelixManager manager = event.getAttribute(AttributeName.helixmanager.name());
  Map<String, Resource> resourceMap =
      event.getAttribute(AttributeName.RESOURCES_TO_REBALANCE.name());
  BaseControllerDataProvider cache = event.getAttribute(AttributeName.ControllerDataProvider.name());
  Map<String, LiveInstance> liveInstanceMap = cache.getLiveInstances();

  if (manager == null || resourceMap == null || messageOutput == null || cache == null
      || liveInstanceMap == null) {
    throw new StageException("Missing attributes in event:" + event
        + ". Requires HelixManager|RESOURCES|MESSAGES_THROTTLE|DataCache|liveInstanceMap");
  }

  HelixDataAccessor dataAccessor = manager.getHelixDataAccessor();
  List<Message> messagesToSend = new ArrayList<>();
  for (String resourceName : resourceMap.keySet()) {
    Resource resource = resourceMap.get(resourceName);
    for (Partition partition : resource.getPartitions()) {
      List<Message> messages = messageOutput.getMessages(resourceName, partition);
      if (messages == null || messages.isEmpty()) {
        messages = Collections.emptyList();
      }
      messagesToSend.addAll(messages);
    }
  }

  List<Message> outputMessages =
      batchMessage(dataAccessor.keyBuilder(), messagesToSend, resourceMap, liveInstanceMap,
          manager.getProperties());

  List<Message> messagesSent = sendMessages(dataAccessor, outputMessages);
  // TODO: Need also count messages from task rebalancer
  if (!(cache instanceof WorkflowControllerDataProvider)) {
    ClusterStatusMonitor clusterStatusMonitor =
        event.getAttribute(AttributeName.clusterStatusMonitor.name());
    if (clusterStatusMonitor != null) {
      clusterStatusMonitor.increaseMessageReceived(outputMessages);
    }
  }
  long cacheStart = System.currentTimeMillis();
  cache.cacheMessages(messagesSent);
  long cacheEnd = System.currentTimeMillis();
  LogUtil.logDebug(logger, _eventId, "Caching messages took " + (cacheEnd - cacheStart) + " ms");
}
 

public static List<String> liveInstances(HelixManager helixManager) {
  HelixDataAccessor helixDataAccessor = helixManager.getHelixDataAccessor();
  PropertyKey liveInstancesKey = helixDataAccessor.keyBuilder().liveInstances();
  return ImmutableList.copyOf(helixDataAccessor.getChildNames(liveInstancesKey));
}
 

@Override
public void onControllerChange(NotificationContext changeContext) {
  logger.info("START: GenericClusterController.onControllerChange() for cluster " + _clusterName);

  requestDataProvidersFullRefresh();

  boolean controllerIsLeader;

  if (changeContext == null || changeContext.getType() == NotificationContext.Type.FINALIZE) {
    logger.info(
        "GenericClusterController.onControllerChange() Cluster change type {} for cluster {}. Disable leadership.",
        changeContext == null ? null : changeContext.getType(), _clusterName);
    controllerIsLeader = false;
  } else {
    // double check if this controller is the leader
    controllerIsLeader = changeContext.getManager().isLeader();
  }

  if (controllerIsLeader) {
    HelixManager manager = changeContext.getManager();
    HelixDataAccessor accessor = manager.getHelixDataAccessor();
    Builder keyBuilder = accessor.keyBuilder();
    PauseSignal pauseSignal = accessor.getProperty(keyBuilder.pause());
    MaintenanceSignal maintenanceSignal = accessor.getProperty(keyBuilder.maintenance());
    _paused = updateControllerState(changeContext, pauseSignal, _paused);
    _inMaintenanceMode =
        updateControllerState(changeContext, maintenanceSignal, _inMaintenanceMode);
    enableClusterStatusMonitor(true);
    _clusterStatusMonitor.setEnabled(!_paused);
    _clusterStatusMonitor.setPaused(_paused);
    _clusterStatusMonitor.setMaintenance(_inMaintenanceMode);
  } else {
    enableClusterStatusMonitor(false);
    // Note that onControllerChange is executed in parallel with the event processing thread. It
    // is possible that the current WAGED rebalancer object is in use for handling callback. So
    // mark the rebalancer invalid only, instead of closing it here.
    // This to-be-closed WAGED rebalancer will be reset later on a later event processing if
    // the controller becomes leader again.
    _rebalancerRef.invalidateRebalancer();
  }

  logger.info("END: GenericClusterController.onControllerChange() for cluster " + _clusterName);
}
 

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

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

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

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

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

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

  Thread.sleep(200);

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

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

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

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

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

/**
 * Creates a setup where the health API can be tested.
 * @param clusterName
 * @param resourceName
 * @param idealStateParams
 * @param partitionReplicaStates maps partitionName to its replicas' states
 * @throws Exception
 */
private void createDummyMapping(String clusterName, String resourceName,
    Map<String, String> idealStateParams, Map<String, List<String>> partitionReplicaStates)
    throws Exception {
  IdealState idealState = new IdealState(resourceName);
  idealState.setMinActiveReplicas(Integer.parseInt(idealStateParams.get("MinActiveReplicas"))); // 2
  idealState.setStateModelDefRef(idealStateParams.get("StateModelDefRef")); // MasterSlave
  idealState.setMaxPartitionsPerInstance(
      Integer.parseInt(idealStateParams.get("MaxPartitionsPerInstance"))); // 3
  idealState.setReplicas(idealStateParams.get("Replicas")); // 3
  idealState.setNumPartitions(Integer.parseInt(idealStateParams.get("NumPartitions"))); // 3
  idealState.enable(false);

  Map<String, List<String>> partitionNames = new LinkedHashMap<>();
  List<String> dummyPrefList = new ArrayList<>();

  for (int i = 0; i < Integer.parseInt(idealStateParams.get("MaxPartitionsPerInstance")); i++) {
    dummyPrefList.add(ANY_INSTANCE);
    partitionNames.put("p" + i, dummyPrefList);
  }
  idealState.getRecord().getListFields().putAll(partitionNames);

  if (!_gSetupTool.getClusterManagementTool().getClusters().contains(clusterName)) {
    _gSetupTool.getClusterManagementTool().addCluster(clusterName);
  }
  _gSetupTool.getClusterManagementTool().setResourceIdealState(clusterName, resourceName,
      idealState);

  // Set ExternalView's replica states for a given parameter map
  ExternalView externalView = new ExternalView(resourceName);

  Map<String, Map<String, String>> mappingCurrent = new LinkedHashMap<>();

  List<String> partitionReplicaStatesList = new ArrayList<>(partitionReplicaStates.keySet());
  for (int k = 0; k < partitionReplicaStatesList.size(); k++) {
    Map<String, String> replicaStatesForPartition = new LinkedHashMap<>();
    List<String> replicaStateList = partitionReplicaStates.get(partitionReplicaStatesList.get(k));
    for (int i = 0; i < replicaStateList.size(); i++) {
      replicaStatesForPartition.put("r" + i, replicaStateList.get(i));
    }
    mappingCurrent.put("p" + k, replicaStatesForPartition);
  }

  externalView.getRecord().getMapFields().putAll(mappingCurrent);

  HelixManager helixManager = HelixManagerFactory.getZKHelixManager(clusterName, "p1",
      InstanceType.ADMINISTRATOR, ZK_ADDR);
  helixManager.connect();
  HelixDataAccessor helixDataAccessor = helixManager.getHelixDataAccessor();
  helixDataAccessor.setProperty(helixDataAccessor.keyBuilder().externalView(resourceName),
      externalView);
  System.out.println("End test :" + TestHelper.getTestMethodName());
}