Java Code Examples for org.apache.helix.HelixDataAccessor#setProperty()

@Test
public void testGroupCommitAddCurrentStateBack() throws InterruptedException {
  HelixDataAccessor accessor = _manager.getHelixDataAccessor();
  Message initMessage = generateMessage("OFFLINE", "ONLINE");
  accessor.setProperty(
      accessor.keyBuilder().message(_participant.getInstanceName(), initMessage.getMsgId()),
      initMessage);
  Assert.assertTrue(waitForMessageProcessed(accessor, initMessage.getMsgId()));
  Message toOffline = generateMessage("ONLINE", "OFFLINE");
  accessor.setProperty(
      accessor.keyBuilder().message(_participant.getInstanceName(), toOffline.getMsgId()),
      toOffline);
  Assert.assertTrue(waitForMessageProcessed(accessor, toOffline.getMsgId()));

  // Consequential 10 messages
  for (int i = 0; i < 10; i++) {
    Message dropped = generateMessage("OFFLINE", "DROPPED");
    accessor.setProperty(
        accessor.keyBuilder().message(_participant.getInstanceName(), dropped.getMsgId()),
        dropped);
    Assert.assertTrue(waitForMessageProcessed(accessor, dropped.getMsgId()));
    Assert.assertFalse(accessor.getBaseDataAccessor()
        .exists(accessor.keyBuilder().currentState(_participant.getInstanceName(),
            _participant.getSessionId(), WorkflowGenerator.DEFAULT_TGT_DB).getPath(), 0));
  }
}
 

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

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

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

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

@Override
public Representation post(Representation entity) {
  try {
    String clusterName = (String) getRequest().getAttributes().get("clusterName");
    String modelName = (String) getRequest().getAttributes().get("modelName");
    ZkClient zkClient =
        (ZkClient) getContext().getAttributes().get(RestAdminApplication.ZKCLIENT);

    JsonParameters jsonParameters = new JsonParameters(entity);
    String command = jsonParameters.getCommand();

    if (command.equalsIgnoreCase(ClusterSetup.addStateModelDef)) {
      ZNRecord newStateModel =
          jsonParameters.getExtraParameter(JsonParameters.NEW_STATE_MODEL_DEF);
      HelixDataAccessor accessor =
          ClusterRepresentationUtil.getClusterDataAccessor(zkClient, clusterName);

      accessor.setProperty(accessor.keyBuilder().stateModelDef(newStateModel.getId()),
          new StateModelDefinition(newStateModel));
    } else {
      throw new HelixException("Unsupported command: " + command + ". Should be one of ["
          + ClusterSetup.addStateModelDef + "]");
    }
    getResponse().setEntity(getStateModelRepresentation(clusterName, modelName));
    getResponse().setStatus(Status.SUCCESS_OK);
  } catch (Exception e) {
    getResponse().setEntity(ClusterRepresentationUtil.getErrorAsJsonStringFromException(e),
        MediaType.APPLICATION_JSON);
    getResponse().setStatus(Status.SUCCESS_OK);
    LOG.error("Error in posting " + entity, e);
  }
  return null;
}
 

private void addStateModels(HelixDataAccessor accessor) {
  StateModelDefinition masterSlave =
      new StateModelDefinition(StateModelConfigGenerator.generateConfigForMasterSlave());
  accessor.setProperty(accessor.keyBuilder().stateModelDef(masterSlave.getId()), masterSlave);
  StateModelDefinition onlineOffline =
      new StateModelDefinition(StateModelConfigGenerator.generateConfigForOnlineOffline());
  accessor.setProperty(accessor.keyBuilder().stateModelDef(onlineOffline.getId()), onlineOffline);
}
 

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

private void sendReply(HelixDataAccessor replyDataAccessor, Message message,
    HelixTaskResult taskResult) {
  if (message.getCorrelationId() != null && !message.getMsgType()
      .equals(MessageType.TASK_REPLY.name())) {
    logger.info("Sending reply for message " + message.getCorrelationId());
    _statusUpdateUtil.logInfo(message, HelixTask.class, "Sending reply", _manager);

    taskResult.getTaskResultMap().put("SUCCESS", "" + taskResult.isSuccess());
    taskResult.getTaskResultMap().put("INTERRUPTED", "" + taskResult.isInterrupted());
    if (!taskResult.isSuccess()) {
      taskResult.getTaskResultMap().put("ERRORINFO", taskResult.getMessage());
    }
    Message replyMessage = Message
        .createReplyMessage(message, _manager.getInstanceName(), taskResult.getTaskResultMap());
    replyMessage.setSrcInstanceType(_manager.getInstanceType());

    Builder keyBuilder = replyDataAccessor.keyBuilder();
    if (message.getSrcInstanceType() == InstanceType.PARTICIPANT) {
      replyDataAccessor
          .setProperty(keyBuilder.message(message.getMsgSrc(), replyMessage.getMsgId()),
              replyMessage);
    } else if (message.getSrcInstanceType() == InstanceType.CONTROLLER) {
      replyDataAccessor
          .setProperty(keyBuilder.controllerMessage(replyMessage.getMsgId()), replyMessage);
    }
    _statusUpdateUtil.logInfo(message, HelixTask.class, String
        .format("1 msg replied to %s in cluster %s.", replyMessage.getTgtName(),
            message.getSrcClusterName() == null ?
                _manager.getClusterName() :
                message.getSrcClusterName()), _manager);
  }
}
 

@Override
public void reportHealthReportMessage(ZNRecord healthCheckInfoUpdate) {
  HelixDataAccessor accessor = _helixManager.getHelixDataAccessor();
  Builder keyBuilder = accessor.keyBuilder();
  if(!accessor.setProperty(keyBuilder.healthReport(_instanceName, healthCheckInfoUpdate.getId()),
      new HealthStat(healthCheckInfoUpdate))) {
    LOG.error("Failed to persist health report to zk!");
  }
}
 

private void createIS(HelixDataAccessor accessor, String resourceId, String stateModelDefRef,
    RebalanceMode rebalanceMode) {
  IdealState idealState = new IdealState(resourceId);
  idealState.setRebalanceMode(rebalanceMode);
  idealState.setStateModelDefRef(stateModelDefRef);
  idealState.setNumPartitions(1);
  idealState.setReplicas("1");
  idealState.getRecord().setListField(resourceId + "_0", ImmutableList.of(PARTICIPANT));
  idealState.getRecord().setMapField(resourceId + "_0", ImmutableMap.of(PARTICIPANT, STATE));
  accessor.setProperty(accessor.keyBuilder().idealStates(resourceId), idealState);
}
 

@Test(dependsOnMethods = "testGetAllMessages")
public void testGetMessagesByStateModelDef() throws IOException {
  System.out.println("Start test :" + TestHelper.getTestMethodName());

  String testInstance = CLUSTER_NAME + "localhost_12926"; //Non-live instance
  String messageId = "msg1";
  Message message = new Message(Message.MessageType.STATE_TRANSITION, messageId);
  message.setStateModelDef("MasterSlave");
  message.setFromState("OFFLINE");
  message.setToState("SLAVE");
  message.setResourceName("testResourceName");
  message.setPartitionName("testResourceName_1");
  message.setTgtName("localhost_3");
  message.setTgtSessionId("session_3");
  HelixDataAccessor helixDataAccessor = new ZKHelixDataAccessor(CLUSTER_NAME, _baseAccessor);
  helixDataAccessor.setProperty(helixDataAccessor.keyBuilder().message(testInstance, messageId),
      message);

  String body =
      new JerseyUriRequestBuilder("clusters/{}/instances/{}/messages?stateModelDef=MasterSlave")
          .isBodyReturnExpected(true).format(CLUSTER_NAME, testInstance).get(this);
  JsonNode node = OBJECT_MAPPER.readTree(body);
  int newMessageCount =
      node.get(PerInstanceAccessor.PerInstanceProperties.total_message_count.name()).getIntValue();

  Assert.assertEquals(newMessageCount, 1);

  body =
      new JerseyUriRequestBuilder("clusters/{}/instances/{}/messages?stateModelDef=LeaderStandBy")
          .isBodyReturnExpected(true).format(CLUSTER_NAME, testInstance).get(this);
  node = OBJECT_MAPPER.readTree(body);
  newMessageCount =
      node.get(PerInstanceAccessor.PerInstanceProperties.total_message_count.name()).getIntValue();

  Assert.assertEquals(newMessageCount, 0);
  System.out.println("End test :" + TestHelper.getTestMethodName());
}
 

@Override
public Representation post(Representation entity) {
  try {
    String clusterName = (String) getRequest().getAttributes().get("clusterName");
    ZkClient zkClient =
        (ZkClient) getContext().getAttributes().get(RestAdminApplication.ZKCLIENT);
    ;

    JsonParameters jsonParameters = new JsonParameters(entity);
    String command = jsonParameters.getCommand();

    if (command.equalsIgnoreCase(ClusterSetup.addStateModelDef)) {
      ZNRecord newStateModel =
          jsonParameters.getExtraParameter(JsonParameters.NEW_STATE_MODEL_DEF);
      HelixDataAccessor accessor =
          ClusterRepresentationUtil.getClusterDataAccessor(zkClient, clusterName);

      accessor.setProperty(accessor.keyBuilder().stateModelDef(newStateModel.getId()),
          new StateModelDefinition(newStateModel));
      getResponse().setEntity(getStateModelsRepresentation());
    } else {
      throw new HelixException("Unsupported command: " + command + ". Should be one of ["
          + ClusterSetup.addStateModelDef + "]");
    }

    getResponse().setStatus(Status.SUCCESS_OK);
  } catch (Exception e) {
    getResponse().setEntity(ClusterRepresentationUtil.getErrorAsJsonStringFromException(e),
        MediaType.APPLICATION_JSON);
    getResponse().setStatus(Status.SUCCESS_OK);
    LOG.error("Error in posting " + entity, e);
  }
  return null;
}
 

@Test(dependsOnMethods = "testInstanceAutoJoin")
public void testAutoRegistrationShouldFailWhenWaitingResponse() throws Exception {
  // Create CloudConfig object and add to config
  CloudConfig.Builder cloudConfigBuilder = new CloudConfig.Builder();
  cloudConfigBuilder.setCloudEnabled(true);
  cloudConfigBuilder.setCloudProvider(CloudProvider.AZURE);
  cloudConfigBuilder.setCloudID("TestID");
  CloudConfig cloudConfig = cloudConfigBuilder.build();

  HelixManager manager = _participants[0];
  HelixDataAccessor accessor = manager.getHelixDataAccessor();

  _gSetupTool.addResourceToCluster(CLUSTER_NAME, db3, 60, "OnlineOffline",
      RebalanceMode.FULL_AUTO + "");
  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, db3, 1);
  String instance3 = "localhost_279700";

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

  manager.getConfigAccessor().set(scope, ZKHelixManager.ALLOW_PARTICIPANT_AUTO_JOIN, "true");
  // Write the CloudConfig to Zookeeper
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  accessor.setProperty(keyBuilder.cloudConfig(), cloudConfig);

  MockParticipantManager autoParticipant =
      new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, instance3);
  autoParticipant.syncStart();

  Assert.assertTrue(null == manager.getHelixDataAccessor()
      .getProperty(accessor.keyBuilder().liveInstance(instance3)));
  try {
    manager.getConfigAccessor().getInstanceConfig(CLUSTER_NAME, instance3);
    fail(
        "Exception should be thrown because the instance cannot be added to the cluster due to the disconnection with Azure endpoint");
  } catch (HelixException e) {

  }
}
 

@BeforeClass
public void beforeClass() throws Exception {
  super.beforeClass();

  String clusterName = TestHelper.getTestClassName();
  int n = 2;

  System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));

  TestHelper.setupCluster(clusterName, ZK_ADDR, 12918, // participant port
      "localhost", // participant name prefix
      "TestDB", // resource name prefix
      2, // resources
      2, // partitions per resource
      n, // number of nodes
      2, // replicas
      "MasterSlave", true); // do rebalance

  ClusterControllerManager controller =
      new ClusterControllerManager(ZK_ADDR, clusterName, "controller_0");
  controller.syncStart();

  // start participants
  MockParticipantManager[] participants = new MockParticipantManager[n];
  for (int i = 0; i < n; i++) {
    String instanceName = "localhost_" + (12918 + i);

    participants[i] = new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
    participants[i].syncStart();
  }

  INSTANCE_0 = participants[0].getInstanceName();
  INSTANCE_1 = participants[1].getInstanceName();

  _manager = HelixManagerFactory
      .getZKHelixManager(clusterName, "admin", InstanceType.ADMINISTRATOR, ZK_ADDR);
  _manager.connect();

  _spectator = HelixManagerFactory
      .getZKHelixManager(clusterName, "spectator", InstanceType.SPECTATOR, ZK_ADDR);
  _spectator.connect();

  // Initialize customized state provider
  _customizedStateProvider_participant0 = CustomizedStateProviderFactory.getInstance()
      .buildCustomizedStateProvider(_manager, participants[0].getInstanceName());
  _customizedStateProvider_participant1 = CustomizedStateProviderFactory.getInstance()
      .buildCustomizedStateProvider(_manager, participants[1].getInstanceName());

  _localCustomizedView = new HashMap<>();
  _routingTableProviderDataSources = new HashSet<>();
  _aggregationEnabledTypes = new HashSet<>();

  List<String> customizedStateTypes = Arrays
      .asList(CustomizedStateType.TYPE_A.name(), CustomizedStateType.TYPE_B.name(),
          CustomizedStateType.TYPE_C.name());

  CustomizedStateConfig.Builder customizedStateConfigBuilder =
      new CustomizedStateConfig.Builder();
  customizedStateConfigBuilder.setAggregationEnabledTypes(customizedStateTypes);
  HelixDataAccessor accessor = _manager.getHelixDataAccessor();
  accessor.setProperty(accessor.keyBuilder().customizedStateConfig(),
      customizedStateConfigBuilder.build());
  _aggregationEnabledTypes.addAll(customizedStateTypes);

  Map<PropertyType, List<String>> dataSource = new HashMap<>();
  dataSource.put(PropertyType.CUSTOMIZEDVIEW, customizedStateTypes);
  _routingTableProvider = new RoutingTableProvider(_spectator, dataSource);
  _routingTableProviderDataSources.addAll(customizedStateTypes);
}
 

@Test
public void test() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  int n = 1;

  System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));

  TestHelper.setupCluster(clusterName, ZK_ADDR, 12918, // participant port
      "localhost", // participant name prefix
      "TestDB", // resource name prefix
      1, // resources
      2, // partitions per resource
      n, // number of nodes
      1, // replicas
      "MasterSlave", true); // do rebalance

  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_gZkClient));
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  BaseDataAccessor<ZNRecord> baseAccessor = accessor.getBaseDataAccessor();

  HelixManager manager = mock(HelixManager.class);
  when(manager.getHelixDataAccessor()).thenReturn(accessor);
  when(manager.getClusterName()).thenReturn(clusterName);

  // run dump task without statusUpdates and errors, should not remove any existing
  // statusUpdate/error paths
  ZKPathDataDumpTask task = new ZKPathDataDumpTask(manager, 0L, 0L, Integer.MAX_VALUE);
  task.run();
  PropertyKey controllerStatusUpdateKey = keyBuilder.controllerTaskStatuses();
  Assert.assertTrue(baseAccessor.exists(controllerStatusUpdateKey.getPath(), 0));
  PropertyKey controllerErrorKey = keyBuilder.controllerTaskErrors();
  Assert.assertTrue(baseAccessor.exists(controllerErrorKey.getPath(), 0));
  PropertyKey statusUpdateKey = keyBuilder.stateTransitionStatus("localhost_12918");
  Assert.assertTrue(baseAccessor.exists(statusUpdateKey.getPath(), 0));
  PropertyKey errorKey = keyBuilder.stateTransitionErrors("localhost_12918");

  // add participant status updates and errors
  statusUpdateKey =
      keyBuilder.stateTransitionStatus("localhost_12918", "session_0", "TestDB0", "TestDB0_0");
  accessor.setProperty(statusUpdateKey, new StatusUpdate(new ZNRecord("statusUpdate")));
  errorKey =
      keyBuilder.stateTransitionError("localhost_12918", "session_0", "TestDB0", "TestDB0_0");
  accessor.setProperty(errorKey, new Error(new ZNRecord("error")));

  // add controller status updates and errors
  controllerStatusUpdateKey = keyBuilder.controllerTaskStatus("session_0", "TestDB");
  accessor.setProperty(controllerStatusUpdateKey,
      new StatusUpdate(new ZNRecord("controllerStatusUpdate")));
  controllerErrorKey = keyBuilder.controllerTaskError("TestDB_error");
  accessor.setProperty(controllerErrorKey, new Error(new ZNRecord("controllerError")));

  // run dump task, should remove existing statusUpdate/error paths
  task.run();
  Assert.assertFalse(baseAccessor.exists(controllerStatusUpdateKey.getPath(), 0));
  Assert.assertFalse(baseAccessor.exists(controllerErrorKey.getPath(), 0));
  Assert.assertFalse(baseAccessor.exists(statusUpdateKey.getPath(), 0));
  Assert.assertFalse(baseAccessor.exists(errorKey.getPath(), 0));

  controllerStatusUpdateKey = keyBuilder.controllerTaskStatuses();
  Assert.assertTrue(baseAccessor.exists(controllerStatusUpdateKey.getPath(), 0));
  controllerErrorKey = keyBuilder.controllerTaskErrors();
  Assert.assertTrue(baseAccessor.exists(controllerErrorKey.getPath(), 0));
  statusUpdateKey = keyBuilder.stateTransitionStatus("localhost_12918");
  Assert.assertTrue(baseAccessor.exists(statusUpdateKey.getPath(), 0));
  errorKey = keyBuilder.stateTransitionErrors("localhost_12918");

  deleteCluster(clusterName);
  System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
}
 

@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 testCapacityReached() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  int n = 1;

  System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));

  TestHelper.setupCluster(clusterName, ZK_ADDR, 12918, // participant port
      "localhost", // participant name prefix
      "TestDB", // resource name prefix
      1, // resources
      2, // partitions per resource
      n, // number of nodes
      1, // replicas
      "MasterSlave", true); // do rebalance

  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_gZkClient));
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  BaseDataAccessor<ZNRecord> baseAccessor = accessor.getBaseDataAccessor();

  HelixManager manager = mock(HelixManager.class);
  when(manager.getHelixDataAccessor()).thenReturn(accessor);
  when(manager.getClusterName()).thenReturn(clusterName);

  // run dump task without statusUpdates and errors, should not remove any existing
  // statusUpdate/error paths
  ZKPathDataDumpTask task = new ZKPathDataDumpTask(manager, Long.MAX_VALUE, Long.MAX_VALUE, 1);
  task.run();
  PropertyKey controllerStatusUpdateKey = keyBuilder.controllerTaskStatuses();
  Assert.assertTrue(baseAccessor.exists(controllerStatusUpdateKey.getPath(), 0));
  PropertyKey controllerErrorKey = keyBuilder.controllerTaskErrors();
  Assert.assertTrue(baseAccessor.exists(controllerErrorKey.getPath(), 0));
  PropertyKey statusUpdateKey = keyBuilder.stateTransitionStatus("localhost_12918");
  Assert.assertTrue(baseAccessor.exists(statusUpdateKey.getPath(), 0));
  PropertyKey errorKey = keyBuilder.stateTransitionErrors("localhost_12918");
  Assert.assertTrue(baseAccessor.exists(errorKey.getPath(), 0));

  // add participant status updates and errors
  statusUpdateKey =
      keyBuilder.stateTransitionStatus("localhost_12918", "session_0", "TestDB0", "TestDB0_0");
  accessor.setProperty(statusUpdateKey, new StatusUpdate(new ZNRecord("statusUpdate")));
  errorKey =
      keyBuilder.stateTransitionError("localhost_12918", "session_0", "TestDB0", "TestDB0_0");
  accessor.setProperty(errorKey, new Error(new ZNRecord("error")));

  // add controller status updates and errors (one of each, should not trigger anything)
  controllerStatusUpdateKey = keyBuilder.controllerTaskStatus("session_0", "TestDB");
  accessor.setProperty(controllerStatusUpdateKey,
      new StatusUpdate(new ZNRecord("controllerStatusUpdate")));
  controllerErrorKey = keyBuilder.controllerTaskError("TestDB_error");
  accessor.setProperty(controllerErrorKey, new Error(new ZNRecord("controllerError")));

  // run dump task, should not remove anything because the threshold is not exceeded
  task.run();
  Assert.assertTrue(baseAccessor.exists(controllerStatusUpdateKey.getPath(), 0));
  Assert.assertTrue(baseAccessor.exists(controllerErrorKey.getPath(), 0));
  Assert.assertTrue(baseAccessor.exists(statusUpdateKey.getPath(), 0));
  Assert.assertTrue(baseAccessor.exists(errorKey.getPath(), 0));

  // add a second set of all status updates and errors
  statusUpdateKey =
      keyBuilder.stateTransitionStatus("localhost_12918", "session_0", "TestDB0", "TestDB0_1");
  accessor.setProperty(statusUpdateKey, new StatusUpdate(new ZNRecord("statusUpdate")));
  errorKey =
      keyBuilder.stateTransitionError("localhost_12918", "session_0", "TestDB0", "TestDB0_1");
  accessor.setProperty(errorKey, new Error(new ZNRecord("error")));
  controllerStatusUpdateKey = keyBuilder.controllerTaskStatus("session_0", "TestDB1");
  accessor.setProperty(controllerStatusUpdateKey,
      new StatusUpdate(new ZNRecord("controllerStatusUpdate")));
  controllerErrorKey = keyBuilder.controllerTaskError("TestDB1_error");
  accessor.setProperty(controllerErrorKey, new Error(new ZNRecord("controllerError")));

  // run dump task, should remove everything since capacities are exceeded
  task.run();
  Assert.assertFalse(baseAccessor.exists(controllerStatusUpdateKey.getPath(), 0));
  Assert.assertFalse(baseAccessor.exists(controllerErrorKey.getPath(), 0));
  Assert.assertFalse(baseAccessor.exists(statusUpdateKey.getPath(), 0));
  Assert.assertFalse(baseAccessor.exists(errorKey.getPath(), 0));
  deleteCluster(clusterName);
}
 

@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 testCurrentStatePathLeakingByAsycRemoval() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  final int n = 3;
  final String zkAddr = ZK_ADDR;
  final int mJobUpdateCnt = 500;

  System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));

  TestHelper.setupCluster(clusterName, zkAddr, 12918, "localhost", "TestDB", 1, // resource
      32, // partitions
      n, // nodes
      2, // replicas
      "MasterSlave", true);

  final ClusterControllerManager controller =
      new ClusterControllerManager(zkAddr, clusterName, "controller_0");
  controller.syncStart();

  MockParticipantManager[] participants = new MockParticipantManager[n];
  for (int i = 0; i < n; i++) {
    String instanceName = "localhost_" + (12918 + i);
    participants[i] = new MockParticipantManager(zkAddr, clusterName, instanceName);
    participants[i].syncStart();
  }

  Boolean result = ClusterStateVerifier.verifyByZkCallback(
      new ClusterStateVerifier.BestPossAndExtViewZkVerifier(zkAddr, clusterName));
  Assert.assertTrue(result);

  ClusterSpectatorManager rpManager = new ClusterSpectatorManager(ZK_ADDR, clusterName, "router");
  rpManager.syncStart();
  RoutingTableProvider rp = new RoutingTableProvider(rpManager, PropertyType.CURRENTSTATES);

  LOG.info("add job");
  MockParticipantManager jobParticipant = participants[0];
  String jobSessionId = jobParticipant.getSessionId();
  HelixDataAccessor jobAccesor = jobParticipant.getHelixDataAccessor();
  PropertyKey.Builder jobKeyBuilder = new PropertyKey.Builder(clusterName);
  PropertyKey db0key =
      jobKeyBuilder.currentState(jobParticipant.getInstanceName(), jobSessionId, "TestDB0");
  CurrentState db0 = jobAccesor.getProperty(db0key);
  PropertyKey jobKey =
      jobKeyBuilder.currentState(jobParticipant.getInstanceName(), jobSessionId, "BackupQueue");
  CurrentState cs = new CurrentState("BackupQueue");
  cs.setSessionId(jobSessionId);
  cs.setStateModelDefRef(db0.getStateModelDefRef());

  LOG.info("add job");
  boolean rtJob = false;
  for (int i = 0; i < mJobUpdateCnt; i++) {
    rtJob = jobAccesor.setProperty(jobKey, cs);
  }

  LOG.info("remove job");
  rtJob = jobParticipant.getZkClient().delete(jobKey.getPath());

  // validate the job watch is not leaked.
  Thread.sleep(5000);

  Map<String, Set<String>> listenersByZkPath = ZkTestHelper.getListenersByZkPath(ZK_ADDR);
  boolean jobKeyExists = listenersByZkPath.keySet().contains(jobKey.getPath());
  Assert.assertFalse(jobKeyExists);

  Map<String, List<String>> rpWatchPaths = ZkTestHelper.getZkWatch(rpManager.getZkClient());
  List<String> existWatches = rpWatchPaths.get("existWatches");
  Assert.assertTrue(existWatches.isEmpty());
}
 

/**
 * Test addResourceWithWeight() and validateResourcesForWagedRebalance() by trying to add a resource with incomplete ResourceConfig.
 */
@Test
public void testAddResourceWithWeightAndValidation()
    throws IOException {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  String mockInstance = "MockInstance";
  String testResourcePrefix = "TestResource";
  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.addCluster(clusterName, true);
  admin.addStateModelDef(clusterName, "MasterSlave", new MasterSlaveSMD());

  // Create a dummy instance
  InstanceConfig instanceConfig = new InstanceConfig(mockInstance);
  Map<String, Integer> mockInstanceCapacity =
      ImmutableMap.of("WCU", 100, "RCU", 100, "STORAGE", 100);
  instanceConfig.setInstanceCapacityMap(mockInstanceCapacity);
  admin.addInstance(clusterName, instanceConfig);
  MockParticipantManager mockParticipantManager =
      new MockParticipantManager(ZK_ADDR, clusterName, mockInstance);
  mockParticipantManager.syncStart();

  IdealState idealState = new IdealState(testResourcePrefix);
  idealState.setNumPartitions(3);
  idealState.setStateModelDefRef("MasterSlave");
  idealState.setRebalanceMode(IdealState.RebalanceMode.FULL_AUTO);

  ResourceConfig resourceConfig = new ResourceConfig(testResourcePrefix);
  // validate
  Map<String, Boolean> validationResult = admin.validateResourcesForWagedRebalance(clusterName,
      Collections.singletonList(testResourcePrefix));
  Assert.assertEquals(validationResult.size(), 1);
  Assert.assertFalse(validationResult.get(testResourcePrefix));
  try {
    admin.addResourceWithWeight(clusterName, idealState, resourceConfig);
    Assert.fail();
  } catch (HelixException e) {
    // OK since resourceConfig is empty
  }

  // Set PARTITION_CAPACITY_MAP
  Map<String, String> capacityDataMap =
      ImmutableMap.of("WCU", "1", "RCU", "2", "STORAGE", "3");
  resourceConfig.getRecord()
      .setMapField(ResourceConfig.ResourceConfigProperty.PARTITION_CAPACITY_MAP.name(),
          Collections.singletonMap(ResourceConfig.DEFAULT_PARTITION_KEY,
              OBJECT_MAPPER.writeValueAsString(capacityDataMap)));

  // validate
  validationResult = admin.validateResourcesForWagedRebalance(clusterName,
      Collections.singletonList(testResourcePrefix));
  Assert.assertEquals(validationResult.size(), 1);
  Assert.assertFalse(validationResult.get(testResourcePrefix));

  // Add the capacity key to ClusterConfig
  HelixDataAccessor dataAccessor = new ZKHelixDataAccessor(clusterName, _baseAccessor);
  PropertyKey.Builder keyBuilder = dataAccessor.keyBuilder();
  ClusterConfig clusterConfig = dataAccessor.getProperty(keyBuilder.clusterConfig());
  clusterConfig.setInstanceCapacityKeys(Arrays.asList("WCU", "RCU", "STORAGE"));
  dataAccessor.setProperty(keyBuilder.clusterConfig(), clusterConfig);

  // Should succeed now
  Assert.assertTrue(admin.addResourceWithWeight(clusterName, idealState, resourceConfig));
  // validate
  validationResult = admin.validateResourcesForWagedRebalance(clusterName,
      Collections.singletonList(testResourcePrefix));
  Assert.assertEquals(validationResult.size(), 1);
  Assert.assertTrue(validationResult.get(testResourcePrefix));
}
 

/**
 * Write an error record to zookeeper to the zookeeper store.
 * @param record
 *          the status update record
 * @param instanceName
 *          the instance name
 * @param updateSubPath
 *          the error update sub path
 * @param updateKey
 *          the error update key
 * @param sessionId
 *          the session id
 * @param accessor
 *          the zookeeper data accessor that writes the status update to zookeeper
 * @param isController
 *          if the error log is for a controller instance or not
 */
void publishErrorRecord(ZNRecord record, String instanceName, String updateSubPath,
    String updateKey, String sessionId, HelixDataAccessor accessor, boolean isController) {
  Builder keyBuilder = accessor.keyBuilder();
  if (isController) {
    // TODO need to fix: ERRORS_CONTROLLER doesn't have a form of
    // ../{sessionId}/{subPath}
    accessor.setProperty(keyBuilder.controllerTaskError(updateSubPath), new Error(record));
  } else {
    accessor.updateProperty(keyBuilder.stateTransitionError(instanceName, sessionId,
        updateSubPath, updateKey), new Error(record));
  }
}
 

/**
 * Set the resource config
 * @param accessor Accessor to Helix configs
 * @param resource The resource name
 * @param resourceConfig The resource config to be set
 * @return True if set successfully, otherwise false
 */
private static boolean setResourceConfig(HelixDataAccessor accessor, String resource,
    ResourceConfig resourceConfig) {
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  return accessor.setProperty(keyBuilder.resourceConfig(resource), resourceConfig);
}