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

private void sendNopMessageInternal() {
  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);
    }
  } catch (Exception e) {
    _logger.error(e.toString());
  }
}
 

@Override
public synchronized void transmitHealthReports() {
  synchronized (_healthReportProviderList) {
    for (HealthReportProvider provider : _healthReportProviderList) {
      try {
        Map<String, String> report = provider.getRecentHealthReport();
        Map<String, Map<String, String>> partitionReport =
            provider.getRecentPartitionHealthReport();
        ZNRecord record = new ZNRecord(provider.getReportName());
        if (report != null) {
          record.setSimpleFields(report);
        }
        if (partitionReport != null) {
          record.setMapFields(partitionReport);
        }
        record.setSimpleField(HealthStat.TIMESTAMP_NAME, "" + System.currentTimeMillis());

        HelixDataAccessor accessor = _helixManager.getHelixDataAccessor();
        Builder keyBuilder = accessor.keyBuilder();
        if (!accessor.setProperty(keyBuilder.healthReport(_instanceName, record.getId()),
            new HealthStat(record))) {
          LOG.error("Failed to persist health report to zk!");
        }

        provider.resetStats();
      } catch (Exception e) {
        LOG.error("fail to transmit health report", e);
      }
    }
  }
}
 

private void verifyJobDeleted(String queueName, String jobName) throws Exception {
  HelixDataAccessor accessor = _manager.getHelixDataAccessor();
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();

  Assert.assertNull(accessor.getProperty(keyBuilder.idealStates(jobName)));
  Assert.assertNull(accessor.getProperty(keyBuilder.resourceConfig(jobName)));
  TaskTestUtil.pollForEmptyJobState(_driver, queueName, jobName);
}
 

private long getTopStateUpdateTime(ExternalView ev, String partition, String state) {
  String topStateHost = ev.getStateMap(partition).entrySet().stream()
      .filter(entry -> entry.getValue().equals(state)).findFirst().get().getKey();
  MockParticipantManager participant = Arrays.stream(_participants)
      .filter(instance -> instance.getInstanceName().equals(topStateHost)).findFirst().get();

  HelixDataAccessor accessor = _controller.getHelixDataAccessor();
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  CurrentState currentState = accessor.getProperty(keyBuilder
      .currentState(participant.getInstanceName(), participant.getSessionId(),
          ev.getResourceName()));
  return currentState.getEndTime(partition);
}
 

@Test
public void testDisableResource() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  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
      10, // partitions per resource
      5, // number of nodes
      3, // replicas
      "MasterSlave", true); // do rebalance
  // disable "TestDB0" resource
  ClusterSetup.processCommandLineArgs(new String[] {
      "--zkSvr", ZK_ADDR, "--enableResource", clusterName, "TestDB0", "false"
  });
  BaseDataAccessor<ZNRecord> baseAccessor = new ZkBaseDataAccessor<ZNRecord>(ZK_ADDR);
  HelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, baseAccessor);
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  IdealState idealState = accessor.getProperty(keyBuilder.idealStates("TestDB0"));
  Assert.assertFalse(idealState.isEnabled());
  // enable "TestDB0" resource
  ClusterSetup.processCommandLineArgs(new String[] {
      "--zkSvr", ZK_ADDR, "--enableResource", clusterName, "TestDB0", "true"
  });
  idealState = accessor.getProperty(keyBuilder.idealStates("TestDB0"));
  Assert.assertTrue(idealState.isEnabled());

  TestHelper.dropCluster(clusterName, _gZkClient);
  System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
}
 

/**
 * Get the customized state for a specified resource
 */
public CustomizedState getCustomizedState(String customizedStateName, String resourceName) {
  HelixDataAccessor accessor = _helixManager.getHelixDataAccessor();
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  return (CustomizedState) accessor
      .getProperty(keyBuilder.customizedState(_instanceName, customizedStateName, resourceName));
}
 

@Override
public void setResourceIdealState(String clusterName, String resourceName,
    IdealState idealState) {
  logger
      .info("Set IdealState for resource {} in cluster {} with new IdealState {}.", resourceName,
          clusterName, idealState == null ? "NULL" : idealState.toString());
  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<ZNRecord>(_zkClient));
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();

  accessor.setProperty(keyBuilder.idealStates(resourceName), idealState);
}
 

@Override
public IdealState getResourceIdealState(String clusterName, String resourceName) {
  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<ZNRecord>(_zkClient));
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();

  return accessor.getProperty(keyBuilder.idealStates(resourceName));
}
 

@Test()
public void testAutoRebalance() throws Exception {
  // kill 1 node
  _participants[0].syncStop();

  boolean result = ClusterStateVerifier
      .verifyByZkCallback(new ExternalViewBalancedVerifier(_gZkClient, CLUSTER_NAME, TEST_DB));
  Assert.assertTrue(result);

  // add 2 nodes
  for (int i = 0; i < 2; i++) {
    String storageNodeName = PARTICIPANT_PREFIX + "_" + (1000 + i);
    _gSetupTool.addInstanceToCluster(CLUSTER_NAME, storageNodeName);

    MockParticipantManager participant =
        new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, storageNodeName.replace(':', '_'));
    _extraParticipants.add(participant);
    participant.syncStart();
  }
  Thread.sleep(100);
  result = ClusterStateVerifier.verifyByPolling(
      new ExternalViewBalancedVerifier(_gZkClient, CLUSTER_NAME, TEST_DB), 10000, 100);
  Assert.assertTrue(result);

  result = ClusterStateVerifier
      .verifyByZkCallback(new ExternalViewBalancedVerifier(_gZkClient, CLUSTER_NAME, db2));
  Assert.assertTrue(result);
  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(CLUSTER_NAME, new ZkBaseDataAccessor<>(_gZkClient));
  Builder keyBuilder = accessor.keyBuilder();
  ExternalView ev = accessor.getProperty(keyBuilder.externalView(db2));
  Set<String> instancesSet = new HashSet<>();
  for (String partitionName : ev.getRecord().getMapFields().keySet()) {
    Map<String, String> assignmentMap = ev.getRecord().getMapField(partitionName);
    instancesSet.addAll(assignmentMap.keySet());
  }
  Assert.assertEquals(instancesSet.size(), 2);
}
 

@Override
public void execute(final ClusterEvent event) throws Exception {
  ResourceControllerDataProvider cache = event.getAttribute(AttributeName.ControllerDataProvider.name());
  ClusterConfig clusterConfig = cache.getClusterConfig();

  if (!clusterConfig.isPersistBestPossibleAssignment() && !clusterConfig
      .isPersistIntermediateAssignment()) {
    return;
  }

  BestPossibleStateOutput bestPossibleAssignment =
      event.getAttribute(AttributeName.BEST_POSSIBLE_STATE.name());

  HelixManager helixManager = event.getAttribute(AttributeName.helixmanager.name());
  HelixDataAccessor accessor = helixManager.getHelixDataAccessor();
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  Map<String, Resource> resourceMap = event.getAttribute(AttributeName.RESOURCES.name());

  for (String resourceId : bestPossibleAssignment.resourceSet()) {
    try {
      persistAssignment(resourceMap.get(resourceId), cache, event, bestPossibleAssignment,
          clusterConfig, accessor, keyBuilder);
    } catch (HelixException ex) {
      LogUtil
          .logError(LOG, _eventId, "Failed to persist assignment for resource " + resourceId, ex);
    }
  }
}
 

/**
 * Case where we have one resource in IdealState
 * @throws Exception
 */
@Test
public void testSimple() throws Exception {
  int nodes = 5;
  List<String> instances = new ArrayList<String>();
  for (int i = 0; i < nodes; i++) {
    instances.add("localhost_" + i);
  }
  int partitions = 10;
  int replicas = 1;
  String resourceName = "testResource";
  ZNRecord record =
      DefaultIdealStateCalculator.calculateIdealState(instances, partitions, replicas,
          resourceName, "MASTER", "SLAVE");
  IdealState idealState = new IdealState(record);
  idealState.setStateModelDefRef("MasterSlave");

  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  Builder keyBuilder = accessor.keyBuilder();
  accessor.setProperty(keyBuilder.idealStates(resourceName), idealState);
  event.addAttribute(AttributeName.ControllerDataProvider.name(),
      new ResourceControllerDataProvider());
  ResourceComputationStage stage = new ResourceComputationStage();
  runStage(event, new ReadClusterDataStage());
  runStage(event, stage);

  Map<String, Resource> resource = event.getAttribute(AttributeName.RESOURCES_TO_REBALANCE.name());
  AssertJUnit.assertEquals(1, resource.size());

  AssertJUnit.assertEquals(resource.keySet().iterator().next(), resourceName);
  AssertJUnit.assertEquals(resource.values().iterator().next().getResourceName(), resourceName);
  AssertJUnit.assertEquals(resource.values().iterator().next().getStateModelDefRef(),
      idealState.getStateModelDefRef());
  AssertJUnit
      .assertEquals(resource.values().iterator().next().getPartitions().size(), partitions);
}
 

private void forwardRelayMessages(HelixDataAccessor accessor, Message message,
    long taskCompletionTime) {
  if (message.hasRelayMessages()) {
    Map<String, Message> relayMessages = message.getRelayMessages();
    Builder keyBuilder = accessor.keyBuilder();

    // Ignore all relay messages if participant's session has changed.
    if (!_manager.getSessionId().equals(message.getTgtSessionId())) {
      logger.info(
          "Session id has been changed, ignore all relay messages attached with " + message
              .getId());
      return;
    }

    for (String instance : relayMessages.keySet()) {
      Message msg = relayMessages.get(instance);
      if (msg.getMsgSubType().equals(MessageType.RELAYED_MESSAGE.name())) {
        msg.setRelayTime(taskCompletionTime);
        if (msg.isExpired()) {
          logger.info(
              "Relay message expired, ignore " + msg.getId() + " to instance " + instance);
          continue;
        }
        PropertyKey msgKey = keyBuilder.message(instance, msg.getId());
        boolean success = accessor.getBaseDataAccessor()
            .create(msgKey.getPath(), msg.getRecord(), AccessOption.PERSISTENT);
        if (!success) {
          logger.warn("Failed to send relay message " + msg.getId() + " to " + instance);
        } else {
          logger.info("Send relay message " + msg.getId() + " to " + instance);
        }
      }
    }
  }
}
 

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 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 testSchedulerMsgUsingQueue() throws Exception {
  // Logger.getRootLogger().setLevel(Level.INFO);
  _factory._results.clear();
  Thread.sleep(2000);
  HelixManager manager = null;
  for (int i = 0; i < NODE_NR; i++) {
    _participants[i].getMessagingService().registerMessageHandlerFactory(
        _factory.getMessageType(), _factory);

    manager = _participants[i]; // _startCMResultMap.get(hostDest)._manager;
  }

  Message schedulerMessage =
      new Message(MessageType.SCHEDULER_MSG + "", UUID.randomUUID().toString());
  schedulerMessage.setTgtSessionId("*");
  schedulerMessage.setTgtName("CONTROLLER");
  // TODO: change it to "ADMIN" ?
  schedulerMessage.setSrcName("CONTROLLER");
  schedulerMessage.getRecord().setSimpleField(
      DefaultSchedulerMessageHandlerFactory.SCHEDULER_TASK_QUEUE, "TestSchedulerMsg");
  // Template for the individual message sent to each participant
  Message msg = new Message(_factory.getMessageType(), "Template");
  msg.setTgtSessionId("*");
  msg.setMsgState(MessageState.NEW);

  // Criteria to send individual messages
  Criteria cr = new Criteria();
  cr.setInstanceName("localhost_%");
  cr.setRecipientInstanceType(InstanceType.PARTICIPANT);
  cr.setSessionSpecific(false);
  cr.setResource("%");
  cr.setPartition("%");

  ObjectMapper mapper = new ObjectMapper();
  SerializationConfig serializationConfig = mapper.getSerializationConfig();
  serializationConfig.set(SerializationConfig.Feature.INDENT_OUTPUT, true);

  StringWriter sw = new StringWriter();
  mapper.writeValue(sw, cr);

  String crString = sw.toString();

  schedulerMessage.getRecord().setSimpleField("Criteria", crString);
  schedulerMessage.getRecord().setMapField("MessageTemplate", msg.getRecord().getSimpleFields());
  schedulerMessage.getRecord().setSimpleField("TIMEOUT", "-1");

  HelixDataAccessor helixDataAccessor = manager.getHelixDataAccessor();
  PropertyKey.Builder keyBuilder = helixDataAccessor.keyBuilder();
  helixDataAccessor.createControllerMessage(schedulerMessage);

  for (int i = 0; i < 30; i++) {
    Thread.sleep(2000);
    if (_PARTITIONS == _factory._results.size()) {
      break;
    }
  }

  Assert.assertEquals(_PARTITIONS, _factory._results.size());
  PropertyKey controllerTaskStatus =
      keyBuilder.controllerTaskStatus(MessageType.SCHEDULER_MSG.name(),
          schedulerMessage.getMsgId());

  int messageResultCount = 0;
  for (int i = 0; i < 10; i++) {
    ZNRecord statusUpdate = helixDataAccessor.getProperty(controllerTaskStatus).getRecord();
    Assert.assertTrue(statusUpdate.getMapField("SentMessageCount").get("MessageCount")
        .equals("" + (_PARTITIONS * 3)));
    for (String key : statusUpdate.getMapFields().keySet()) {
      if (key.startsWith("MessageResult ")) {
        messageResultCount++;
      }
    }
    if (messageResultCount == _PARTITIONS * 3) {
      break;
    } else {
      Thread.sleep(2000);
    }
  }
  Assert.assertEquals(messageResultCount, _PARTITIONS * 3);
  int count = 0;
  for (Set<String> val : _factory._results.values()) {
    count += val.size();
  }
  Assert.assertEquals(count, _PARTITIONS * 3);

}
 

@Test
public void testNamedQueue() throws Exception {
  String queueName = TestHelper.getTestMethodName();

  // Create a queue
  JobQueue queue = new JobQueue.Builder(queueName).build();
  _driver.createQueue(queue);

  // Enqueue jobs
  Set<String> master = Sets.newHashSet("MASTER");
  Set<String> slave = Sets.newHashSet("SLAVE");
  JobConfig.Builder job1 = new JobConfig.Builder().setCommand(MockTask.TASK_COMMAND)
      .setTargetResource(WorkflowGenerator.DEFAULT_TGT_DB).setTargetPartitionStates(master);
  JobConfig.Builder job2 = new JobConfig.Builder().setCommand(MockTask.TASK_COMMAND)
      .setTargetResource(WorkflowGenerator.DEFAULT_TGT_DB).setTargetPartitionStates(slave);
  _driver.enqueueJob(queueName, "masterJob", job1);
  _driver.enqueueJob(queueName, "slaveJob", job2);

  // Ensure successful completion
  String namespacedJob1 = queueName + "_masterJob";
  String namespacedJob2 = queueName + "_slaveJob";
  _driver.pollForJobState(queueName, namespacedJob1, TaskState.COMPLETED);
  _driver.pollForJobState(queueName, namespacedJob2, TaskState.COMPLETED);
  JobContext masterJobContext = _driver.getJobContext(namespacedJob1);
  JobContext slaveJobContext = _driver.getJobContext(namespacedJob2);

  // Ensure correct ordering
  long job1Finish = masterJobContext.getFinishTime();
  long job2Start = slaveJobContext.getStartTime();
  Assert.assertTrue(job2Start >= job1Finish);

  // Flush queue and check cleanup
  _driver.flushQueue(queueName);
  HelixDataAccessor accessor = _manager.getHelixDataAccessor();
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  Assert.assertNull(accessor.getProperty(keyBuilder.idealStates(namespacedJob1)));
  Assert.assertNull(accessor.getProperty(keyBuilder.resourceConfig(namespacedJob1)));
  Assert.assertNull(accessor.getProperty(keyBuilder.idealStates(namespacedJob2)));
  Assert.assertNull(accessor.getProperty(keyBuilder.resourceConfig(namespacedJob2)));
  WorkflowConfig workflowCfg = _driver.getWorkflowConfig(queueName);
  JobDag dag = workflowCfg.getJobDag();
  Assert.assertFalse(dag.getAllNodes().contains(namespacedJob1));
  Assert.assertFalse(dag.getAllNodes().contains(namespacedJob2));
  Assert.assertFalse(dag.getChildrenToParents().containsKey(namespacedJob1));
  Assert.assertFalse(dag.getChildrenToParents().containsKey(namespacedJob2));
  Assert.assertFalse(dag.getParentsToChildren().containsKey(namespacedJob1));
  Assert.assertFalse(dag.getParentsToChildren().containsKey(namespacedJob2));
}
 

public static StatusUpdateContents getStatusUpdateContents(HelixDataAccessor accessor,
    String instance, String resourceGroup, String sessionID, String partition) {
  Builder keyBuilder = accessor.keyBuilder();

  List<ZNRecord> instances =
      HelixProperty.convertToList(accessor.getChildValues(keyBuilder.instanceConfigs(), true));
  List<ZNRecord> partitionRecords = new ArrayList<ZNRecord>();
  for (ZNRecord znRecord : instances) {
    String instanceName = znRecord.getId();
    if (!instanceName.equals(instance)) {
      continue;
    }

    List<String> sessions = accessor.getChildNames(keyBuilder.sessions(instanceName));
    for (String session : sessions) {
      if (sessionID != null && !session.equals(sessionID)) {
        continue;
      }

      List<String> resourceGroups =
          accessor.getChildNames(keyBuilder.stateTransitionStatus(instanceName, session));
      for (String resourceGroupName : resourceGroups) {
        if (!resourceGroupName.equals(resourceGroup)) {
          continue;
        }

        List<String> partitionStrings =
            accessor.getChildNames(keyBuilder.stateTransitionStatus(instanceName, session,
                resourceGroupName));

        for (String partitionString : partitionStrings) {
          ZNRecord partitionRecord =
              accessor.getProperty(
                  keyBuilder.stateTransitionStatus(instanceName, session, resourceGroupName,
                      partitionString)).getRecord();
          if (!partitionString.equals(partition)) {
            continue;
          }
          partitionRecords.add(partitionRecord);
        }
      }
    }
  }

  return new StatusUpdateContents(getSortedTransitions(partitionRecords),
      getTaskMessages(partitionRecords));
}
 

@Test
public void testParticipantFlapping() throws Exception {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  final HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_gZkClient));
  final PropertyKey.Builder keyBuilder = accessor.keyBuilder();

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

  System.setProperty(SystemPropertyKeys.MAX_DISCONNECT_THRESHOLD,
      Integer.toString(_disconnectThreshold));

  try {
    TestHelper.setupCluster(clusterName, ZK_ADDR, 12918, // participant port
        "localhost", // participant name prefix
        "TestDB", // resource name prefix
        1, // resources
        32, // partitions per resource
        1, // number of nodes
        1, // replicas
        "MasterSlave", false);

    final String instanceName = "localhost_12918";
    MockParticipantManager participant =
        new MockParticipantManager(ZK_ADDR, clusterName, instanceName);
    participant.syncStart();

    final ZkClient client = (ZkClient) participant.getZkClient();
    final ZkStateCountListener listener = new ZkStateCountListener();
    client.subscribeStateChanges(listener);

    final AtomicInteger expectDisconnectCnt = new AtomicInteger(0);
    final int n = _disconnectThreshold;
    for (int i = 0; i < _disconnectThreshold; i++) {
      String oldSessionId = ZkTestHelper.getSessionId(client);
      ZkTestHelper.simulateZkStateReconnected(client);
      expectDisconnectCnt.incrementAndGet();
      // wait until we get invoked by zk state change to disconnected
      TestHelper.verify(() -> listener.count == expectDisconnectCnt.get(), 30 * 1000);

      String newSessionId = ZkTestHelper.getSessionId(client);
      Assert.assertEquals(newSessionId, oldSessionId);
    }
    client.unsubscribeStateChanges(listener);
    // make sure participant is NOT disconnected
    LiveInstance liveInstance = accessor.getProperty(keyBuilder.liveInstance(instanceName));
    Assert.assertNotNull(liveInstance, "Live-instance should exist after " + n + " disconnects");

    // trigger flapping
    ZkTestHelper.simulateZkStateReconnected(client);
    // wait until we get invoked by zk state change to disconnected
    boolean success = TestHelper.verify(client::getShutdownTrigger, 30 * 1000);

    Assert.assertTrue(success,
        "The " + (n + 1) + "th disconnect event should trigger ZkHelixManager#disonnect");

    // make sure participant is disconnected
    success = TestHelper.verify(() -> {
      LiveInstance liveInstance1 = accessor.getProperty(keyBuilder.liveInstance(instanceName));
      return liveInstance1 == null;
    }, 3 * 1000);
    Assert.assertTrue(success, "Live-instance should be gone after " + (n + 1) + " disconnects");
    participant.syncStop();
  } finally {
    System.clearProperty(SystemPropertyKeys.MAX_DISCONNECT_THRESHOLD);
  }

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

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

void checkDuplicatedMessages(List<Message> messages) {
  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  PropertyKey path = keyBuilder.currentStates(manager.getInstanceName(), manager.getSessionId());
  Map<String, CurrentState> currentStateMap = accessor.getChildValuesMap(path, true);

  Set<String> seenPartitions = new HashSet<>();
  for (Message message : messages) {
    if (message.getMsgType().equals(Message.MessageType.STATE_TRANSITION.name())) {
      String resource = message.getResourceName();
      String partition = message.getPartitionName();

      //System.err.println(message.getMsgId());
      String key = resource + "-" + partition;
      if (seenPartitions.contains(key)) {
        //System.err.println("Duplicated message received for " + resource + ":" + partition);
        duplicatedMessages++;
      }
      seenPartitions.add(key);

      String toState = message.getToState();
      String state = null;
      if (currentStateMap.containsKey(resource)) {
        CurrentState currentState = currentStateMap.get(resource);
        state = currentState.getState(partition);
      }

      if (toState.equals(state) && message.getMsgState() == Message.MessageState.NEW) {
        //            logger.error(
        //                "Extra message: " + message.getMsgId() + ", Partition is already in target state "
        //                    + toState + " for " + resource + ":" + partition);
        extraStateTransition++;
      }

      String messageTarget =
          getMessageTarget(message.getResourceName(), message.getPartitionName());

      if (message.getMsgState() == Message.MessageState.NEW &&
          _messageTaskMap.containsKey(messageTarget)) {
        String taskId = _messageTaskMap.get(messageTarget);
        MessageTaskInfo messageTaskInfo = _taskMap.get(taskId);
        Message existingMsg = messageTaskInfo.getTask().getMessage();
        if (existingMsg.getMsgId() != message.getMsgId())
          //            logger.error("Duplicated message In Progress: " + message.getMsgId()
          //                    + ", state transition in progress with message " + existingMsg.getMsgId()
          //                    + " to " + toState + " for " + resource + ":" + partition);
          duplicatedMessagesInProgress ++;
      }
    }
  }
}