org.apache.helix.model.CurrentState Java Examples

private void verifyP2PDisabled() {
  ResourceControllerDataProvider dataCache = new ResourceControllerDataProvider(CLUSTER_NAME);
  dataCache.refresh(_accessor);
  Map<String, LiveInstance> liveInstanceMap = dataCache.getLiveInstances();

  for (LiveInstance instance : liveInstanceMap.values()) {
    Map<String, CurrentState> currentStateMap =
        dataCache.getCurrentState(instance.getInstanceName(), instance.getEphemeralOwner());
    Assert.assertNotNull(currentStateMap);
    for (CurrentState currentState : currentStateMap.values()) {
      for (String partition : currentState.getPartitionStateMap().keySet()) {
        String state = currentState.getState(partition);
        if (state.equalsIgnoreCase("MASTER")) {
          String triggerHost = currentState.getTriggerHost(partition);
          Assert.assertEquals(triggerHost, _controllerName,
              state + " of " + partition + " on " + instance.getInstanceName()
                  + " was triggered by " + triggerHost);
        }
      }
    }
  }
}
 

@GET @Path("resources/{resourceName}")
public Response getResourceOnInstance(@PathParam("clusterId") String clusterId,
    @PathParam("instanceName") String instanceName,
    @PathParam("resourceName") String resourceName) throws IOException {
  HelixDataAccessor accessor = getDataAccssor(clusterId);
  List<String> sessionIds = accessor.getChildNames(accessor.keyBuilder().sessions(instanceName));
  if (sessionIds == null || sessionIds.size() == 0) {
    return notFound();
  }

  // Only get resource list from current session id
  String currentSessionId = sessionIds.get(0);
  CurrentState resourceCurrentState = accessor.getProperty(
      accessor.keyBuilder().currentState(instanceName, currentSessionId, resourceName));
  if (resourceCurrentState != null) {
    return JSONRepresentation(resourceCurrentState.getRecord());
  }

  return notFound();
}
 

private void verifyP2PEnabled(long startTime) {
  ResourceControllerDataProvider dataCache = new ResourceControllerDataProvider(CLUSTER_NAME);
  dataCache.refresh(_accessor);
  Map<String, LiveInstance> liveInstanceMap = dataCache.getLiveInstances();

  for (LiveInstance instance : liveInstanceMap.values()) {
    Map<String, CurrentState> currentStateMap =
        dataCache.getCurrentState(instance.getInstanceName(), instance.getEphemeralOwner());
    Assert.assertNotNull(currentStateMap);
    for (CurrentState currentState : currentStateMap.values()) {
      for (String partition : currentState.getPartitionStateMap().keySet()) {
        String state = currentState.getState(partition);
        long start = currentState.getStartTime(partition);
        if (state.equalsIgnoreCase("MASTER") && start > startTime) {
          String triggerHost = currentState.getTriggerHost(partition);
          if (!triggerHost.equals(_controllerName)) {
            p2pTrigged ++;
          }
          total ++;
        }
      }
    }
  }
}
 

/**
 * Assert externalView and currentState for each participant are empty
 * @param clusterName
 * @param db
 * @param participants
 */
private void assertEmptyCSandEV(String clusterName, String db,
    MockParticipantManager[] participants) throws Exception {
  HelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_gZkClient));
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  boolean isExternalViewNull = TestHelper.verify(() -> {
    ExternalView externalView = accessor.getProperty(keyBuilder.externalView(db));
    return (externalView == null);
  }, TestHelper.WAIT_DURATION);
  Assert.assertTrue(isExternalViewNull);

  for (MockParticipantManager participant : participants) {
    String instanceName = participant.getInstanceName();
    String sessionId = participant.getSessionId();
    boolean isCurrentStateNull = TestHelper.verify(() -> {
      CurrentState currentState = accessor.getProperty(keyBuilder.currentState(instanceName, sessionId, db));
      return (currentState == null);
    }, TestHelper.WAIT_DURATION);
    Assert.assertTrue(isCurrentStateNull);
  }
}
 

@Test public void testAddedFieldsInCurrentState() {
  String instanceName = PARTICIPANT_PREFIX + "_" + _startPort;
  HelixDataAccessor accessor = _manager.getHelixDataAccessor();
  LiveInstance liveInstance =
      accessor.getProperty(accessor.keyBuilder().liveInstance(instanceName));
  CurrentState currentState = accessor.getProperty(accessor.keyBuilder()
      .currentState(instanceName, liveInstance.getEphemeralOwner(), WorkflowGenerator.DEFAULT_TGT_DB));
  // Test start time should happen after test start time
  Assert.assertTrue(
      currentState.getStartTime(WorkflowGenerator.DEFAULT_TGT_DB + "_0") >= _testStartTime);

  // Test end time is always larger than start time
  Assert.assertTrue(
      currentState.getEndTime(WorkflowGenerator.DEFAULT_TGT_DB + "_0") >= currentState
          .getStartTime(WorkflowGenerator.DEFAULT_TGT_DB + "_0"));

  // Final state is MASTER, so SLAVE will be the previous state
  Assert.assertEquals(currentState.getPreviousState(WorkflowGenerator.DEFAULT_TGT_DB + "_0"),
      "SLAVE");
}
 

@Test
public void TestHasErrorPartitions_false() {
  String sessionId = "sessionId";
  String resource = "db";
  Mock mock = new Mock();
  LiveInstance liveInstance = new LiveInstance(TEST_INSTANCE);
  liveInstance.setSessionId(sessionId);
  doReturn(liveInstance).when(mock.dataAccessor)
      .getProperty(argThat(new PropertyKeyArgument(PropertyType.LIVEINSTANCES)));
  doReturn(ImmutableList.of(resource)).when(mock.dataAccessor)
      .getChildNames(argThat(new PropertyKeyArgument(PropertyType.CURRENTSTATES)));
  CurrentState currentState = mock(CurrentState.class);
  when(currentState.getPartitionStateMap()).thenReturn(ImmutableMap.of("db0", "Master"));
  doReturn(currentState).when(mock.dataAccessor)
      .getProperty(argThat(new PropertyKeyArgument(PropertyType.CURRENTSTATES)));

  Assert.assertFalse(
      InstanceValidationUtil.hasErrorPartitions(mock.dataAccessor, TEST_CLUSTER, TEST_INSTANCE));
}
 

@Test
public void TestHasErrorPartitions_true() {
  String sessionId = "sessionId";
  String resource = "db";
  Mock mock = new Mock();
  LiveInstance liveInstance = new LiveInstance(TEST_INSTANCE);
  liveInstance.setSessionId(sessionId);
  doReturn(liveInstance).when(mock.dataAccessor)
      .getProperty(argThat(new PropertyKeyArgument(PropertyType.LIVEINSTANCES)));
  doReturn(ImmutableList.of(resource)).when(mock.dataAccessor)
      .getChildNames(argThat(new PropertyKeyArgument(PropertyType.CURRENTSTATES)));
  CurrentState currentState = mock(CurrentState.class);
  when(currentState.getPartitionStateMap()).thenReturn(ImmutableMap.of("db0", "ERROR"));
  doReturn(currentState).when(mock.dataAccessor)
      .getProperty(argThat(new PropertyKeyArgument(PropertyType.CURRENTSTATES)));

  Assert.assertTrue(
      InstanceValidationUtil.hasErrorPartitions(mock.dataAccessor, TEST_CLUSTER, TEST_INSTANCE));
}
 

@Test
public void TestHasResourceAssigned_fail() {
  String sessionId = "sessionId";
  String resource = "db";
  Mock mock = new Mock();
  LiveInstance liveInstance = new LiveInstance(TEST_INSTANCE);
  liveInstance.setSessionId(sessionId);
  doReturn(liveInstance).when(mock.dataAccessor)
      .getProperty(argThat(new PropertyKeyArgument(PropertyType.LIVEINSTANCES)));
  doReturn(ImmutableList.of(resource)).when(mock.dataAccessor)
      .getChildNames(argThat(new PropertyKeyArgument(PropertyType.CURRENTSTATES)));
  CurrentState currentState = mock(CurrentState.class);
  when(currentState.getPartitionStateMap()).thenReturn(Collections.<String, String> emptyMap());
  doReturn(currentState).when(mock.dataAccessor)
      .getProperty(argThat(new PropertyKeyArgument(PropertyType.CURRENTSTATES)));

  Assert.assertFalse(
      InstanceValidationUtil.hasResourceAssigned(mock.dataAccessor, TEST_CLUSTER, TEST_INSTANCE));
}
 

@Test
public void TestHasResourceAssigned_success() {
  String sessionId = "sessionId";
  String resource = "db";
  Mock mock = new Mock();
  LiveInstance liveInstance = new LiveInstance(TEST_INSTANCE);
  liveInstance.setSessionId(sessionId);
  doReturn(liveInstance).when(mock.dataAccessor)
      .getProperty(argThat(new PropertyKeyArgument(PropertyType.LIVEINSTANCES)));
  doReturn(ImmutableList.of(resource)).when(mock.dataAccessor)
      .getChildNames(argThat(new PropertyKeyArgument(PropertyType.CURRENTSTATES)));
  CurrentState currentState = mock(CurrentState.class);
  when(currentState.getPartitionStateMap()).thenReturn(ImmutableMap.of("db0", "master"));
  doReturn(currentState).when(mock.dataAccessor)
      .getProperty(argThat(new PropertyKeyArgument(PropertyType.CURRENTSTATES)));

  Assert.assertTrue(
      InstanceValidationUtil.hasResourceAssigned(mock.dataAccessor, TEST_CLUSTER, TEST_INSTANCE));
}
 

/**
 * Generate the replica objects according to the provider information.
 */
protected Set<AssignableReplica> generateReplicas(ResourceControllerDataProvider dataProvider) {
  // Create assignable replica based on the current state.
  Map<String, CurrentState> currentStatemap =
      dataProvider.getCurrentState(_testInstanceId, _sessionId);
  Set<AssignableReplica> assignmentSet = new HashSet<>();
  for (CurrentState cs : currentStatemap.values()) {
    ResourceConfig resourceConfig = dataProvider.getResourceConfig(cs.getResourceName());
    // Construct one AssignableReplica for each partition in the current state.
    cs.getPartitionStateMap().entrySet().stream()
        .forEach(entry -> assignmentSet
            .add(new AssignableReplica(dataProvider.getClusterConfig(), resourceConfig,
                entry.getKey(), entry.getValue(), entry.getValue().equals("MASTER") ? 1 : 2)));
  }
  return assignmentSet;
}
 

/**
 * Request a state change for a specific task.
 *
 * @param accessor  connected Helix data accessor
 * @param instance  the instance serving the task
 * @param sessionId the current session of the instance
 * @param resource  the job name
 * @param partition the task partition name
 * @param state     the requested state
 * @return true if the request was persisted, false otherwise
 */
private static boolean setRequestedState(HelixDataAccessor accessor, String instance,
    String sessionId, String resource, String partition, TaskPartitionState state) {
  LOG.debug(
      String.format("Requesting a state transition to %s for partition %s.", state, partition));
  try {
    PropertyKey.Builder keyBuilder = accessor.keyBuilder();
    PropertyKey key = keyBuilder.currentState(instance, sessionId, resource);
    CurrentState currStateDelta = new CurrentState(resource);
    currStateDelta.setRequestedState(partition, state.name());

    return accessor.updateProperty(key, currStateDelta);
  } catch (Exception e) {
    LOG.error(String
        .format("Error when requesting a state transition to %s for partition %s.", state,
            partition), e);
    return false;
  }
}
 

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

@Override
@PreFetch(enabled = false)
public void onStateChange(String instanceName, List<CurrentState> statesInfo,
    NotificationContext changeContext) {
  if (_isBlocking) {
    try {
      newEventHandlingCount.acquire();
    } catch (InterruptedException e) {
      throw new HelixException("Failed to acquire handling lock for testing.");
    }
  }
  super.onStateChange(instanceName, statesInfo, changeContext);
}
 

@Override
protected synchronized void refreshCurrentState(
    Map<String, Map<String, Map<String, CurrentState>>> currentStateMap,
    Collection<InstanceConfig> instanceConfigs, Collection<LiveInstance> liveInstances,
    String referenceKey) {
  super.refreshCurrentState(currentStateMap, instanceConfigs, liveInstances, "Test");
  _refreshCount++;
  if (DEBUG) {
    print();
  }
}
 

protected void refreshSnapshot(Map<PropertyKey, CurrentState> newStateCache,
    Map<PropertyKey, CurrentState> participantStateCache, Set<PropertyKey> reloadedKeys) {
  if (_initialized) {
    _snapshot = new CurrentStateSnapshot(newStateCache, participantStateCache, reloadedKeys);
  } else {
    _snapshot = new CurrentStateSnapshot(newStateCache);
    _initialized = true;
  }
}
 

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

void addCurStateUpdate(Message subMessage, PropertyKey key, CurrentState delta) {
  String parentMid = subMessage.getAttribute(Attributes.PARENT_MSG_ID);
  GroupMessageInfo info = _groupMsgMap.get(parentMid);
  if (info != null) {
    info._curStateUpdateList.add(new CurrentStateUpdate(key, delta));
  }

}
 

protected void refreshCurrentState(Map<String, Map<String, Map<String, CurrentState>>> currentStateMap,
    Collection<InstanceConfig> instanceConfigs, Collection<LiveInstance> liveInstances,
    String referenceKey) {
  long startTime = System.currentTimeMillis();
  RoutingTable newRoutingTable =
      new RoutingTable(currentStateMap, instanceConfigs, liveInstances);
  resetRoutingTableAndNotify(startTime, newRoutingTable, referenceKey);
}
 

public CurStateCarryOverUpdater(String curSessionId, String initState, CurrentState lastCurState) {
  if (curSessionId == null || initState == null || lastCurState == null) {
    throw new IllegalArgumentException(
        "missing curSessionId|initState|lastCurState for carry-over");
  }
  _curSessionId = curSessionId;
  _initState = initState;
  _lastCurState = lastCurState;
}
 

public HelixStateTransitionHandler(StateModelFactory<? extends StateModel> stateModelFactory,
    StateModel stateModel, Message message, NotificationContext context,
    CurrentState currentStateDelta) {
  super(message, context);
  _stateModel = stateModel;
  _statusUpdateUtil = new StatusUpdateUtil();
  _transitionMethodFinder = new StateModelParser();
  _currentStateDelta = currentStateDelta;
  _manager = _notificationContext.getManager();
  _stateModelFactory = stateModelFactory;
}
 

private void verifyP2PMessage(String dbName, String instance, String expectedState,
    String expectedTriggerHost, double expectedRatio) {
  ResourceControllerDataProvider dataCache = new ResourceControllerDataProvider(CLUSTER_NAME);
  dataCache.refresh(_accessor);

  Map<String, LiveInstance> liveInstanceMap = dataCache.getLiveInstances();
  LiveInstance liveInstance = liveInstanceMap.get(instance);

  Map<String, CurrentState> currentStateMap =
      dataCache.getCurrentState(instance, liveInstance.getEphemeralOwner());
  Assert.assertNotNull(currentStateMap);
  CurrentState currentState = currentStateMap.get(dbName);
  Assert.assertNotNull(currentState);
  Assert.assertEquals(currentState.getPartitionStateMap().size(), PARTITION_NUMBER);

  int total = 0;
  int expectedHost = 0;
  for (String partition : currentState.getPartitionStateMap().keySet()) {
    String state = currentState.getState(partition);
    Assert.assertEquals(state, expectedState,
        dbName + " Partition " + partition + "'s state is different as expected!");
    String triggerHost = currentState.getTriggerHost(partition);
    if (triggerHost.equals(expectedTriggerHost)) {
      expectedHost++;
    }
    total++;
  }

  double ratio = ((double) expectedHost) / ((double) total);
  Assert.assertTrue(ratio >= expectedRatio,
      String.format(
          "Only %d out of %d percent transitions to Master were triggered by expected host!",
          expectedHost, total));
}
 

private void refresh(Map<String, Map<String, Map<String, CurrentState>>> currentStateMap) {
  Map<String, InstanceConfig> instanceConfigMap = new HashMap<>();
  if (currentStateMap != null && !currentStateMap.isEmpty()) {
    for (InstanceConfig config : _instanceConfigs) {
      instanceConfigMap.put(config.getId(), config);
    }
    for (LiveInstance liveInstance : _liveInstances) {
      String instanceName = liveInstance.getInstanceName();
      String sessionId = liveInstance.getEphemeralOwner();
      InstanceConfig instanceConfig = instanceConfigMap.get(instanceName);
      if (instanceConfig == null) {
        logger.warn(
            "Participant {} is not found with proper configuration information. It might already be removed from the cluster. "
                + "Skip recording partition assignments that are related to this instance.",
            instanceName);
        continue;
      }

      Map<String, CurrentState> currentStates = Collections.emptyMap();
      if (currentStateMap.containsKey(instanceName)
          && currentStateMap.get(instanceName).containsKey(sessionId)) {
        currentStates = currentStateMap.get(instanceName).get(sessionId);
      }

      for (CurrentState currentState : currentStates.values()) {
        String resourceName = currentState.getResourceName();
        Map<String, String> stateMap = currentState.getPartitionStateMap();

        for (String partitionName : stateMap.keySet()) {
          String state = stateMap.get(partitionName);
          addEntry(resourceName, partitionName, state, instanceConfig);
        }
      }
    }
  }
}
 

private Map<String, CurrentState> generateCurrentStateMap(
    Map<String, CurrentStateInfo> currentStateRawData) {
  Map<String, CurrentState> currentStateMap = new HashMap<String, CurrentState>();
  for (String instanceName : currentStateRawData.keySet()) {
    CurrentStateInfo info = currentStateRawData.get(instanceName);
    CurrentState currentState = new CurrentState(TEST_RESOURCE);
    currentState.setSessionId(SESSION_PREFIX + instanceName.split("_")[1]);
    currentState.setState(PARTITION, info.currentState);
    currentState.setPreviousState(PARTITION, info.previousState);
    currentState.setStartTime(PARTITION, info.startTime);
    currentState.setEndTime(PARTITION, info.endTime);
    currentStateMap.put(instanceName, currentState);
  }
  return currentStateMap;
}
 

protected void setupCurrentStates(Map<String, CurrentState> currentStates) {
  Builder keyBuilder = accessor.keyBuilder();
  for (String instanceName : currentStates.keySet()) {
    accessor.setProperty(keyBuilder
        .currentState(instanceName, currentStates.get(instanceName).getSessionId(),
            currentStates.get(instanceName).getResourceName()), currentStates.get(instanceName));
  }
}
 

private void setCurrentState(String clusterName, String instance, String resourceGroupName,
    String resourceKey, String sessionId, String state, boolean updateTimestamp) {
  ZKHelixDataAccessor accessor =
      new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(_gZkClient));
  Builder keyBuilder = accessor.keyBuilder();

  CurrentState curState = new CurrentState(resourceGroupName);
  curState.setState(resourceKey, state);
  curState.setSessionId(sessionId);
  curState.setStateModelDefRef("MasterSlave");
  if (updateTimestamp) {
    curState.setEndTime(resourceKey, System.currentTimeMillis());
  }
  accessor.setProperty(keyBuilder.currentState(instance, sessionId, resourceGroupName), curState);
}
 

public static boolean verifyEmptyCurStateAndExtView(String clusterName, String resourceName,
    Set<String> instanceNames, String zkAddr) {
  HelixZkClient zkClient = SharedZkClientFactory.getInstance()
      .buildZkClient(new HelixZkClient.ZkConnectionConfig(zkAddr));
  zkClient.setZkSerializer(new ZNRecordSerializer());

  try {
    ZKHelixDataAccessor accessor =
        new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<ZNRecord>(zkClient));
    Builder keyBuilder = accessor.keyBuilder();

    for (String instanceName : instanceNames) {
      List<String> sessionIds = accessor.getChildNames(keyBuilder.sessions(instanceName));

      for (String sessionId : sessionIds) {
        CurrentState curState =
            accessor.getProperty(keyBuilder.currentState(instanceName, sessionId, resourceName));

        if (curState != null && curState.getRecord().getMapFields().size() != 0) {
          return false;
        }
      }

      ExternalView extView = accessor.getProperty(keyBuilder.externalView(resourceName));

      if (extView != null && extView.getRecord().getMapFields().size() != 0) {
        return false;
      }
    }

    return true;
  } finally {
    zkClient.close();
  }
}
 

@Test()
public void testInvocation() throws Exception {
  HelixTaskExecutor executor = new HelixTaskExecutor();
  System.out.println("START TestCMTaskHandler.testInvocation()");
  Message message = new Message(MessageType.STATE_TRANSITION, "Some unique id");

  message.setSrcName("cm-instance-0");
  message.setTgtSessionId("1234");
  message.setFromState("Offline");
  message.setToState("Slave");
  message.setPartitionName("TestDB_0");
  message.setMsgId("Some unique message id");
  message.setResourceName("TestDB");
  message.setTgtName("localhost");
  message.setStateModelDef("MasterSlave");
  message.setStateModelFactoryName(HelixConstants.DEFAULT_STATE_MODEL_FACTORY);
  MockMasterSlaveStateModel stateModel = new MockMasterSlaveStateModel();
  NotificationContext context;
  MockManager manager = new MockManager("clusterName");
  HelixDataAccessor accessor = manager.getHelixDataAccessor();
  StateModelDefinition stateModelDef =
      new StateModelDefinition(StateModelConfigGenerator.generateConfigForMasterSlave());
  Builder keyBuilder = accessor.keyBuilder();
  accessor.setProperty(keyBuilder.stateModelDef("MasterSlave"), stateModelDef);

  context = new NotificationContext(manager);
  CurrentState currentStateDelta = new CurrentState("TestDB");
  currentStateDelta.setState("TestDB_0", "OFFLINE");

  HelixStateTransitionHandler stHandler =
      new HelixStateTransitionHandler(null, stateModel, message, context, currentStateDelta);
  HelixTask handler;
  handler = new HelixTask(message, context, stHandler, executor);
  handler.call();
  AssertJUnit.assertTrue(stateModel.stateModelInvoked);
  System.out.println("END TestCMTaskHandler.testInvocation() at "
      + new Date(System.currentTimeMillis()));
}
 

@Override
public boolean verify() {
  BaseDataAccessor<ZNRecord> baseAccessor = new ZkBaseDataAccessor<ZNRecord>(_zkClient);
  HelixDataAccessor accessor = new ZKHelixDataAccessor(_clusterName, baseAccessor);
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  ExternalView externalView = accessor.getProperty(keyBuilder.externalView(_resourceName));

  // verify external view empty
  if (externalView != null) {
    for (String partition : externalView.getPartitionSet()) {
      Map<String, String> stateMap = externalView.getStateMap(partition);
      if (stateMap != null && !stateMap.isEmpty()) {
        LOG.error("External view not empty for " + partition);
        return false;
      }
    }
  }

  // verify current state empty
  List<String> liveParticipants = accessor.getChildNames(keyBuilder.liveInstances());
  for (String participant : liveParticipants) {
    List<String> sessionIds = accessor.getChildNames(keyBuilder.sessions(participant));
    for (String sessionId : sessionIds) {
      CurrentState currentState =
          accessor.getProperty(keyBuilder.currentState(participant, sessionId, _resourceName));
      Map<String, String> partitionStateMap = currentState.getPartitionStateMap();
      if (partitionStateMap != null && !partitionStateMap.isEmpty()) {
        LOG.error("Current state not empty for " + participant);
        return false;
      }
    }
  }
  return true;
}
 

@Test(description = "test getNewCurrentStateEndTimes")
public void testGetNewCurrentStateEndTimes() {
  String instance1 = "instance1";
  String session1 = "session1";
  String resource1 = "resource1";

  String partition1 = "partition1";
  String partition2 = "partition2";

  PropertyKey key = new PropertyKey.Builder("cluster").currentState(instance1, session1, resource1);

  CurrentState nxtState = new CurrentState(resource1);
  // partition 1, expect to record in endTimesMap
  nxtState.setState(partition1, "SLAVE");
  nxtState.setEndTime(partition1, 200);
  // partition 2, expect to not record in endTimeMap. This is fixing current 1.4T observed timestamp issue
  nxtState.setState(partition2, "MASTER");

  Map<PropertyKey, CurrentState> currentStateMap = new HashMap<>();
  Map<PropertyKey, CurrentState> nextStateMap = new HashMap<>();
  nextStateMap.put(key, nxtState);

  Set<PropertyKey> updateKeys = new HashSet<>();
  updateKeys.add(key);

  CurrentStateSnapshot snapshot = new CurrentStateSnapshot(nextStateMap, currentStateMap, updateKeys);

  Map<PropertyKey, Map<String, Long>> endTimesMap = snapshot.getNewCurrentStateEndTimes();

  Assert.assertEquals(endTimesMap.size(), 1);
  Assert.assertTrue(endTimesMap.get(key).get(partition1) == 200);
}
 

/**
 * Returns the current state for the given resource, or {@code null} if instance is not live or current state does
 * not exist.
 */
@Nullable
@Override
protected CurrentState getState(String resourceName) {
  PropertyKey.Builder keyBuilder = _helixDataAccessor.keyBuilder();
  LiveInstance liveInstance = _helixDataAccessor.getProperty(keyBuilder.liveInstance(_instanceName));
  if (liveInstance == null) {
    return null;
  } else {
    String sessionId = liveInstance.getSessionId();
    return _helixDataAccessor.getProperty(keyBuilder.currentState(_instanceName, sessionId, resourceName));
  }
}