Java Code Examples for org.apache.helix.model.LiveInstance#setSessionId()

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

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

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

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

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

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

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

private void prepare(String controllerVersion, String participantVersion,
    String minSupportedParticipantVersion) {
  List<String> instances =
      Arrays.asList("localhost_0", "localhost_1", "localhost_2", "localhost_3", "localhost_4");
  int partitions = 10;
  int replicas = 1;

  // set ideal state
  String resourceName = "testResource";
  ZNRecord record =
      DefaultIdealStateCalculator.calculateIdealState(instances, partitions, replicas,
          resourceName, "MASTER", "SLAVE");
  IdealState idealState = new IdealState(record);
  idealState.setStateModelDefRef("MasterSlave");

  Builder keyBuilder = accessor.keyBuilder();
  accessor.setProperty(keyBuilder.idealStates(resourceName), idealState);

  // set live instances
  record = new ZNRecord("localhost_0");
  if (participantVersion != null) {
    record.setSimpleField(LiveInstanceProperty.HELIX_VERSION.toString(), participantVersion);
  }
  LiveInstance liveInstance = new LiveInstance(record);
  liveInstance.setSessionId("session_0");
  accessor.setProperty(keyBuilder.liveInstance("localhost_0"), liveInstance);

  if (controllerVersion != null) {
    ((MockManager) manager).setVersion(controllerVersion);
  }

  if (minSupportedParticipantVersion != null) {
    manager.getProperties().getProperties()
        .put("minimum_supported_version.participant", minSupportedParticipantVersion);
  }
  event.addAttribute(AttributeName.helixmanager.name(), manager);
  event.addAttribute(AttributeName.ControllerDataProvider.name(),
      new ResourceControllerDataProvider());
  runStage(event, new ReadClusterDataStage());
}
 

protected List<String> setupLiveInstances(int numLiveInstances) {
  List<String> instances = new ArrayList<>();
  for (int i = 0; i < numLiveInstances; i++) {
    LiveInstance liveInstance = new LiveInstance(HOSTNAME_PREFIX + i);
    liveInstance.setSessionId(SESSION_PREFIX + i);

    Builder keyBuilder = accessor.keyBuilder();
    accessor.setProperty(keyBuilder.liveInstance(HOSTNAME_PREFIX + i), liveInstance);
    instances.add(liveInstance.getInstanceName());
  }

  return instances;
}
 

private Map<String, Map<String, String>> getMapping(final Map<String, List<String>> listResult) {
  final Map<String, Map<String, String>> mapResult = new HashMap<String, Map<String, String>>();
  ResourceControllerDataProvider cache = new ResourceControllerDataProvider();
  MockAccessor accessor = new MockAccessor();
  Builder keyBuilder = accessor.keyBuilder();
  ClusterConfig clusterConfig = new ClusterConfig("TestCluster");
  accessor.setProperty(keyBuilder.clusterConfig(), clusterConfig);
  for (String node : _liveNodes) {
    LiveInstance liveInstance = new LiveInstance(node);
    liveInstance.setSessionId("testSession");
    accessor.setProperty(keyBuilder.liveInstance(node), liveInstance);
  }
  cache.refresh(accessor);

  IdealState is = new IdealState("resource");
  for (String partition : _partitions) {
    List<String> preferenceList = listResult.get(partition);
    Map<String, String> currentStateMap = _currentMapping.get(partition);
    Set<String> disabled = Collections.emptySet();
    Partition p = new Partition(partition);
    CurrentStateOutput currentStateOutput = new CurrentStateOutput();
    if (currentStateMap != null) {
      for (String instance : currentStateMap.keySet()) {
        currentStateOutput
            .setCurrentState("resource", p, instance, currentStateMap.get(instance));
      }
    }
    Map<String, String> assignment = new AutoRebalancer()
        .computeBestPossibleStateForPartition(cache.getLiveInstances().keySet(), _stateModelDef,
            preferenceList, currentStateOutput, disabled, is, clusterConfig, p,
            MonitoredAbnormalResolver.DUMMY_STATE_RESOLVER);
    mapResult.put(partition, assignment);
  }
  return mapResult;
}
 

protected List<LiveInstance> setupLiveInstances(String clusterName, int[] liveInstances) {
  HelixZkClient.ZkClientConfig clientConfig = new HelixZkClient.ZkClientConfig();
  clientConfig.setZkSerializer(new ZNRecordSerializer());

  List<LiveInstance> result = new ArrayList<>();

  for (int i = 0; i < liveInstances.length; i++) {
    String instance = "localhost_" + liveInstances[i];

    _liveInstanceOwners.putIfAbsent(clusterName, new HashMap<>());
    Map<String, HelixZkClient> clientMap = _liveInstanceOwners.get(clusterName);
    clientMap.putIfAbsent(instance, DedicatedZkClientFactory.getInstance()
        .buildZkClient(new HelixZkClient.ZkConnectionConfig(ZK_ADDR), clientConfig));
    HelixZkClient client = clientMap.get(instance);

        ZKHelixDataAccessor accessor =
        new ZKHelixDataAccessor(clusterName, new ZkBaseDataAccessor<>(client));
    Builder keyBuilder = accessor.keyBuilder();

    LiveInstance liveInstance = new LiveInstance(instance);
    // Keep setting the session id in the deprecated field for ensure the same behavior as a real participant.
    // Note the participant is doing so for backward compatibility.
    liveInstance.setSessionId(Long.toHexString(client.getSessionId()));
    // Please refer to the version requirement here: helix-core/src/main/resources/cluster-manager-version.properties
    // Ensuring version compatibility can avoid the warning message during test.
    liveInstance.setHelixVersion("0.4");
    accessor.setProperty(keyBuilder.liveInstance(instance), liveInstance);
    result.add(accessor.getProperty(keyBuilder.liveInstance(instance)));
  }
  return result;
}
 

protected LiveInstance createMockLiveInstance(String instanceId) {
  LiveInstance testLiveInstance = new LiveInstance(instanceId);
  testLiveInstance.setSessionId(_sessionId);
  return testLiveInstance;
}
 

@Test()
public void testInvocation() throws Exception {

  HelixManager testHelixManager =
      HelixManagerFactory.getZKHelixManager(clusterName, "localhost_8900",
          InstanceType.PARTICIPANT, ZK_ADDR);
  testHelixManager.connect();

  TestZKCallback test = new TestZKCallback();

  TestZKCallback.TestCallbackListener testListener = test.new TestCallbackListener();

  testHelixManager.addMessageListener(testListener, "localhost_8900");
  testHelixManager.addCurrentStateChangeListener(testListener, "localhost_8900",
      testHelixManager.getSessionId());
  testHelixManager.addCustomizedStateRootChangeListener(testListener, "localhost_8900");
  testHelixManager.addConfigChangeListener(testListener);
  testHelixManager.addIdealStateChangeListener(testListener);
  testHelixManager.addExternalViewChangeListener(testListener);
  testHelixManager.addLiveInstanceChangeListener(testListener);
  // Initial add listener should trigger the first execution of the
  // listener callbacks
  AssertJUnit.assertTrue(testListener.configChangeReceived
      & testListener.currentStateChangeReceived & testListener.externalViewChangeReceived
      & testListener.idealStateChangeReceived & testListener.liveInstanceChangeReceived
      & testListener.messageChangeReceived);

  testListener.Reset();
  HelixDataAccessor accessor = testHelixManager.getHelixDataAccessor();
  Builder keyBuilder = accessor.keyBuilder();

  ExternalView extView = new ExternalView("db-12345");
  accessor.setProperty(keyBuilder.externalView("db-12345"), extView);
  Thread.sleep(100);
  AssertJUnit.assertTrue(testListener.externalViewChangeReceived);
  testListener.Reset();

  CurrentState curState = new CurrentState("db-12345");
  curState.setSessionId("sessionId");
  curState.setStateModelDefRef("StateModelDef");
  accessor.setProperty(keyBuilder.currentState("localhost_8900", testHelixManager.getSessionId(),
      curState.getId()), curState);
  Thread.sleep(100);
  AssertJUnit.assertTrue(testListener.currentStateChangeReceived);
  testListener.Reset();

  IdealState idealState = new IdealState("db-1234");
  idealState.setNumPartitions(400);
  idealState.setReplicas(Integer.toString(2));
  idealState.setStateModelDefRef("StateModeldef");
  accessor.setProperty(keyBuilder.idealStates("db-1234"), idealState);
  Thread.sleep(100);
  AssertJUnit.assertTrue(testListener.idealStateChangeReceived);
  testListener.Reset();

  // dummyRecord = new ZNRecord("db-12345");
  // dataAccessor.setProperty(PropertyType.IDEALSTATES, idealState, "db-12345"
  // );
  // Thread.sleep(100);
  // AssertJUnit.assertTrue(testListener.idealStateChangeReceived);
  // testListener.Reset();

  // dummyRecord = new ZNRecord("localhost:8900");
  // List<ZNRecord> recList = new ArrayList<ZNRecord>();
  // recList.add(dummyRecord);

  testListener.Reset();
  Message message = new Message(MessageType.STATE_TRANSITION, UUID.randomUUID().toString());
  message.setTgtSessionId("*");
  message.setResourceName("testResource");
  message.setPartitionName("testPartitionKey");
  message.setStateModelDef("MasterSlave");
  message.setToState("toState");
  message.setFromState("fromState");
  message.setTgtName("testTarget");
  message.setStateModelFactoryName(HelixConstants.DEFAULT_STATE_MODEL_FACTORY);

  accessor.setProperty(keyBuilder.message("localhost_8900", message.getId()), message);
  Thread.sleep(500);
  AssertJUnit.assertTrue(testListener.messageChangeReceived);

  // dummyRecord = new ZNRecord("localhost_9801");
  LiveInstance liveInstance = new LiveInstance("localhost_9801");
  liveInstance.setSessionId(UUID.randomUUID().toString());
  liveInstance.setHelixVersion(UUID.randomUUID().toString());
  accessor.setProperty(keyBuilder.liveInstance("localhost_9801"), liveInstance);
  Thread.sleep(500);
  AssertJUnit.assertTrue(testListener.liveInstanceChangeReceived);
  testListener.Reset();

  // dataAccessor.setNodeConfigs(recList); Thread.sleep(100);
  // AssertJUnit.assertTrue(testListener.configChangeReceived);
  // testListener.Reset();

  accessor.removeProperty(keyBuilder.liveInstance("localhost_8900"));
  accessor.removeProperty(keyBuilder.liveInstance("localhost_9801"));
}