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

public static Set<TopicPartition> getUnassignedPartitions(HelixManager helixManager) {
  Set<TopicPartition> unassignedPartitions = new HashSet<TopicPartition>();
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  String helixClusterName = helixManager.getClusterName();
  for (String topic : helixAdmin.getResourcesInCluster(helixClusterName)) {
    IdealState is = helixAdmin.getResourceIdealState(helixClusterName, topic);
    int numPartitions = is.getNumPartitions();
    for (int partition = 0; partition < numPartitions; ++partition) {
      if (is.getInstanceSet(Integer.toString(partition)).isEmpty()) {
        TopicPartition tpi = new TopicPartition(topic, partition);
        unassignedPartitions.add(tpi);
      }
    }
  }
  return unassignedPartitions;
}
 

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

protected void addFakeBrokerInstanceToAutoJoinHelixCluster(String instanceId, boolean isSingleTenant)
    throws Exception {
  HelixManager helixManager = HelixManagerFactory
      .getZKHelixManager(getHelixClusterName(), instanceId, InstanceType.PARTICIPANT, ZkStarter.DEFAULT_ZK_STR);
  helixManager.getStateMachineEngine()
      .registerStateModelFactory(FakeBrokerResourceOnlineOfflineStateModelFactory.STATE_MODEL_DEF,
          FakeBrokerResourceOnlineOfflineStateModelFactory.FACTORY_INSTANCE);
  helixManager.connect();
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  if (isSingleTenant) {
    helixAdmin.addInstanceTag(getHelixClusterName(), instanceId, TagNameUtils.getBrokerTagForTenant(null));
  } else {
    helixAdmin.addInstanceTag(getHelixClusterName(), instanceId, UNTAGGED_BROKER_INSTANCE);
  }
  _fakeInstanceHelixManagers.add(helixManager);
}
 

protected void addFakeServerInstanceToAutoJoinHelixCluster(String instanceId, boolean isSingleTenant, int adminPort)
    throws Exception {
  HelixManager helixManager = HelixManagerFactory
      .getZKHelixManager(getHelixClusterName(), instanceId, InstanceType.PARTICIPANT, ZkStarter.DEFAULT_ZK_STR);
  helixManager.getStateMachineEngine()
      .registerStateModelFactory(FakeSegmentOnlineOfflineStateModelFactory.STATE_MODEL_DEF,
          FakeSegmentOnlineOfflineStateModelFactory.FACTORY_INSTANCE);
  helixManager.connect();
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  if (isSingleTenant) {
    helixAdmin.addInstanceTag(getHelixClusterName(), instanceId, TagNameUtils.getOfflineTagForTenant(null));
    helixAdmin.addInstanceTag(getHelixClusterName(), instanceId, TagNameUtils.getRealtimeTagForTenant(null));
  } else {
    helixAdmin.addInstanceTag(getHelixClusterName(), instanceId, UNTAGGED_SERVER_INSTANCE);
  }
  HelixConfigScope configScope =
      new HelixConfigScopeBuilder(HelixConfigScope.ConfigScopeProperty.PARTICIPANT, getHelixClusterName())
          .forParticipant(instanceId).build();
  helixAdmin.setConfig(configScope, Collections.singletonMap(ADMIN_PORT_KEY, Integer.toString(adminPort)));
  _fakeInstanceHelixManagers.add(helixManager);
}
 

private void assignIdealStates(HelixManager helixManager,
    Map<String, IdealState> idealStatesFromAssignment) {
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  String helixClusterName = helixManager.getClusterName();
  for (String topic : idealStatesFromAssignment.keySet()) {
    IdealState idealState = idealStatesFromAssignment.get(topic);
    helixAdmin.setResourceIdealState(helixClusterName, topic, idealState);
  }
}
 

/**
 * From IdealStates.
 *
 * @return InstanceToNumTopicPartitionMap
 */
public static Map<String, Set<TopicPartition>> getInstanceToTopicPartitionsMap(
    HelixManager helixManager,
    Map<String, KafkaBrokerTopicObserver> clusterToObserverMap) {
  Map<String, Set<TopicPartition>> instanceToNumTopicPartitionMap = new HashMap<>();
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  String helixClusterName = helixManager.getClusterName();
  for (String topic : helixAdmin.getResourcesInCluster(helixClusterName)) {
    IdealState is = helixAdmin.getResourceIdealState(helixClusterName, topic);
    for (String partition : is.getPartitionSet()) {
      TopicPartition tpi;
      if (partition.startsWith("@")) {
        if (clusterToObserverMap != null) {
          TopicPartition topicParition = clusterToObserverMap.get(getSrcFromRoute(partition))
              .getTopicPartitionWithRefresh(topic);
          int trueNumPartition = topicParition != null ? topicParition.getPartition() : -1;
          tpi = new TopicPartition(topic, trueNumPartition, partition);
        } else {
          tpi = new TopicPartition(topic, -1, partition);
        }
      } else {
        // route
        tpi = new TopicPartition(topic, Integer.parseInt(partition));
      }
      for (String instance : is.getInstanceSet(partition)) {
        instanceToNumTopicPartitionMap.putIfAbsent(instance, new HashSet<>());
        instanceToNumTopicPartitionMap.get(instance).add(tpi);
      }
    }
  }
  return instanceToNumTopicPartitionMap;
}
 

IdealStateMatchServiceStatusCallback(HelixManager helixManager, String clusterName, String instanceName,
    List<String> resourcesToMonitor, double minResourcesStartPercent) {
  _clusterName = clusterName;
  _instanceName = instanceName;
  _helixAdmin = helixManager.getClusterManagmentTool();
  _helixDataAccessor = helixManager.getHelixDataAccessor();

  _resourcesToMonitor = new HashSet<>(resourcesToMonitor);
  _numTotalResourcesToMonitor = _resourcesToMonitor.size();

  _minResourcesStartCount = (int) Math.ceil(minResourcesStartPercent * _numTotalResourcesToMonitor / 100);
  LOGGER.info("Monitoring {} resources: {} for start up of instance {}", _numTotalResourcesToMonitor,
      getResourceListAsString(), _instanceName);
}
 

public TaskDriver(HelixManager manager) {
  this(manager.getClusterManagmentTool(), manager.getHelixDataAccessor(),
      manager.getHelixPropertyStore(), manager.getClusterName());
}
 

@Test
public void testPinotOnlyController() {
  ControllerConf firstPinotOnlyControllerConfig = getDefaultControllerConfiguration();
  firstPinotOnlyControllerConfig.setControllerMode(ControllerConf.ControllerMode.PINOT_ONLY);
  ControllerStarter firstPinotOnlyController = getControllerStarter(firstPinotOnlyControllerConfig);

  // Starting Pinot-only controller without Helix controller should fail
  try {
    firstPinotOnlyController.start();
    Assert.fail("Starting Pinot-only controller without Helix controller should fail");
  } catch (Exception e) {
    // Expected
  }

  // Start a Helix-only controller
  ControllerConf helixOnlyControllerConfig = getDefaultControllerConfiguration();
  helixOnlyControllerConfig.setControllerMode(ControllerConf.ControllerMode.HELIX_ONLY);
  helixOnlyControllerConfig.setControllerPort(Integer.toString(DEFAULT_CONTROLLER_PORT + 1));
  ControllerStarter helixOnlyController = new ControllerStarter(helixOnlyControllerConfig);
  helixOnlyController.start();
  HelixManager helixControllerManager = helixOnlyController.getHelixControllerManager();
  HelixAdmin helixAdmin = helixControllerManager.getClusterManagmentTool();
  TestUtils.waitForCondition(aVoid -> helixControllerManager.isConnected(), TIMEOUT_IN_MS,
      "Failed to start the Helix-only controller");

  // Start the first Pinot-only controller
  firstPinotOnlyController.start();
  PinotHelixResourceManager helixResourceManager = firstPinotOnlyController.getHelixResourceManager();
  TestUtils.waitForCondition(aVoid -> helixResourceManager.getHelixZkManager().isConnected(), TIMEOUT_IN_MS,
      "Failed to start the first Pinot-only controller");
  // The first Pinot-only controller should be the MASTER for all partitions
  checkInstanceState(helixAdmin);

  // Start the second Pinot-only controller
  ControllerConf secondPinotOnlyControllerConfig = getDefaultControllerConfiguration();
  secondPinotOnlyControllerConfig.setControllerMode(ControllerConf.ControllerMode.PINOT_ONLY);
  secondPinotOnlyControllerConfig.setControllerPort(Integer.toString(DEFAULT_CONTROLLER_PORT + 2));
  ControllerStarter secondPinotOnlyController = getControllerStarter(secondPinotOnlyControllerConfig);
  secondPinotOnlyController.start();
  TestUtils.waitForCondition(
      aVoid -> secondPinotOnlyController.getHelixResourceManager().getHelixZkManager().isConnected(), TIMEOUT_IN_MS,
      "Failed to start the second Pinot-only controller");
  // There should still be only one MASTER instance for each partition
  checkInstanceState(helixAdmin);

  // Stop the second Pinot-only controller, and there should still be only one MASTER instance for each partition
  secondPinotOnlyController.stop();
  checkInstanceState(helixAdmin);

  // Stop the first Pinot-only controller, and there should be no partition in the external view
  firstPinotOnlyController.stop();
  TestUtils.waitForCondition(aVoid -> {
    ExternalView leadControllerResourceExternalView =
        helixAdmin.getResourceExternalView(getHelixClusterName(), Helix.LEAD_CONTROLLER_RESOURCE_NAME);
    return leadControllerResourceExternalView.getPartitionSet().isEmpty();
  }, TIMEOUT_IN_MS, "Without live instance, there should be no partition in the external view");

  // Stop the Helix-only controller
  helixOnlyController.stop();
}
 

/**
 * Tests setReplicaSealedState method for {@link HelixParticipant}
 * @throws IOException
 */
@Test
public void testGetAndSetReplicaSealedState() throws IOException {
  //setup HelixParticipant and dependencies
  String partitionIdStr = "somePartitionId";
  String partitionIdStr2 = "someOtherPartitionId";
  ReplicaId replicaId = createMockAmbryReplica(partitionIdStr);
  ReplicaId replicaId2 = createMockAmbryReplica(partitionIdStr2);
  String hostname = "localhost";
  int port = 2200;
  String instanceName = ClusterMapUtils.getInstanceName(hostname, port);
  ClusterMapConfig clusterMapConfig = new ClusterMapConfig(new VerifiableProperties(props));
  HelixParticipant helixParticipant =
      new HelixParticipant(clusterMapConfig, helixManagerFactory, new MetricRegistry(),
          getDefaultZkConnectStr(clusterMapConfig), true);
  helixParticipant.participate(Collections.emptyList());
  HelixManager helixManager = helixManagerFactory.getZKHelixManager(null, null, null, null);
  HelixAdmin helixAdmin = helixManager.getClusterManagmentTool();
  InstanceConfig instanceConfig = new InstanceConfig("someInstanceId");
  helixAdmin.setInstanceConfig(clusterName, instanceName, instanceConfig);

  //Make sure the current sealedReplicas list is empty
  List<String> sealedReplicas = helixParticipant.getSealedReplicas();
  assertEquals("sealedReplicas should be empty", Collections.emptyList(), sealedReplicas);

  String listName = "sealedReplicas";

  //Check that invoking setReplicaSealedState with a non-AmbryReplica ReplicaId throws an IllegalArgumentException
  ReplicaId notAmbryReplica = createMockNotAmbryReplica(partitionIdStr);
  try {
    helixParticipant.setReplicaSealedState(notAmbryReplica, true);
    fail("Expected an IllegalArgumentException here");
  } catch (IllegalArgumentException e) {
    //Expected exception
  }

  //Check that invoking setReplicaSealedState adds the partition to the list of sealed replicas
  helixParticipant.setReplicaSealedState(replicaId, true);
  sealedReplicas = helixParticipant.getSealedReplicas();
  listIsExpectedSize(sealedReplicas, 1, listName);
  assertTrue(sealedReplicas.contains(partitionIdStr));

  //Seal another replicaId
  helixParticipant.setReplicaSealedState(replicaId2, true);
  sealedReplicas = helixParticipant.getSealedReplicas();
  listIsExpectedSize(sealedReplicas, 2, listName);
  assertTrue(sealedReplicas.contains(partitionIdStr2));
  assertTrue(sealedReplicas.contains(partitionIdStr));

  //Check that sealed replica list doesn't take duplicates (and that dups are detected by partitionId comparison, not
  //replicaId object comparison
  ReplicaId dup = createMockAmbryReplica(partitionIdStr);
  helixParticipant.setReplicaSealedState(dup, true);
  helixParticipant.setReplicaSealedState(replicaId2, true);
  sealedReplicas = helixParticipant.getSealedReplicas();
  listIsExpectedSize(sealedReplicas, 2, listName);
  assertTrue(sealedReplicas.contains(partitionIdStr2));
  assertTrue(sealedReplicas.contains(partitionIdStr));

  //Check that invoking setReplicaSealedState with isSealed == false removes partition from list of sealed replicas
  helixParticipant.setReplicaSealedState(replicaId, false);
  sealedReplicas = helixParticipant.getSealedReplicas();
  listIsExpectedSize(sealedReplicas, 1, listName);
  assertTrue(sealedReplicas.contains(partitionIdStr2));
  assertFalse(sealedReplicas.contains(partitionIdStr));

  //Removing a replicaId that's already been removed doesn't hurt anything
  helixParticipant.setReplicaSealedState(replicaId, false);
  sealedReplicas = helixParticipant.getSealedReplicas();
  listIsExpectedSize(sealedReplicas, 1, listName);

  //Removing all replicas yields expected behavior (and removal works by partitionId, not replicaId itself)
  dup = createMockAmbryReplica(partitionIdStr2);
  helixParticipant.setReplicaSealedState(dup, false);
  sealedReplicas = helixParticipant.getSealedReplicas();
  listIsExpectedSize(sealedReplicas, 0, listName);
}