import org.apache.helix.HelixAdmin; //导入依赖的package包/类
/**
 * From IdealStates.
 *
 * @return InstanceToNumTopicPartitionMap
 */
public static Map<String, Set<TopicPartition>> getInstanceToTopicPartitionsMap(
    HelixManager helixManager) {
  Map<String, Set<TopicPartition>> instanceToNumTopicPartitionMap =
      new HashMap<String, Set<TopicPartition>>();
  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 = new TopicPartition(topic, Integer.parseInt(partition));
      for (String instance : is.getInstanceSet(partition)) {
        if (!instanceToNumTopicPartitionMap.containsKey(instance)) {
          instanceToNumTopicPartitionMap.put(instance, new HashSet<TopicPartition>());
        }
        instanceToNumTopicPartitionMap.get(instance).add(tpi);
      }
    }
  }
  return instanceToNumTopicPartitionMap;
}
 

import org.apache.helix.HelixAdmin; //导入依赖的package包/类
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;
}
 

import org.apache.helix.HelixAdmin; //导入依赖的package包/类
/**
 * From IdealStates.
 * @param helixManager
 * @return InstanceToNumTopicPartitionMap
 */
public static Map<String, Set<TopicPartition>> getInstanceToTopicPartitionsMap(
    HelixManager helixManager) {
  Map<String, Set<TopicPartition>> instanceToNumTopicPartitionMap =
      new HashMap<String, Set<TopicPartition>>();
  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 = new TopicPartition(topic, Integer.parseInt(partition));
      for (String instance : is.getInstanceSet(partition)) {
        if (!instanceToNumTopicPartitionMap.containsKey(instance)) {
          instanceToNumTopicPartitionMap.put(instance, new HashSet<TopicPartition>());
        }
        instanceToNumTopicPartitionMap.get(instance).add(tpi);
      }
    }
  }
  return instanceToNumTopicPartitionMap;
}
 

import org.apache.helix.HelixAdmin; //导入依赖的package包/类
public TerrapinControllerServiceImpl(PropertiesConfiguration configuration,
                                     ZooKeeperManager zkManager,
                                     DFSClient hdfsClient,
                                     HelixAdmin helixAdmin,
                                     String clusterName) {
  this.configuration = configuration;
  this.zkManager = zkManager;
  this.hdfsClient = hdfsClient;
  this.helixAdmin = helixAdmin;
  this.clusterName = clusterName;

  ExecutorService threadPool = new ThreadPoolExecutor(100,
      100,
      0,
      TimeUnit.SECONDS,
      new LinkedBlockingDeque<Runnable>(1000),
      new ThreadFactoryBuilder().setDaemon(false)
                    .setNameFormat("controller-pool-%d")
                    .build());
 this.futurePool = new ExecutorServiceFuturePool(threadPool);
}
 

import org.apache.helix.HelixAdmin; //导入依赖的package包/类
@Override
public void onLiveInstanceChange(List<LiveInstance> list, NotificationContext notificationContext) {
  List<InetSocketAddress> addresses = new ArrayList<>();
  HelixAdmin helixAdmin = helixConnection.createClusterManagementTool();
  for (LiveInstance liveInstance : list) {
    InstanceConfig instanceConfig = helixAdmin.getInstanceConfig(
        clusterName,
        liveInstance.getInstanceName());
    InetSocketAddress address = new InetSocketAddress(
        instanceConfig.getHostName(),
        Integer.valueOf(instanceConfig.getPort()));
    addresses.add(address);
  }

  services.set(addresses);
}
 

import org.apache.helix.HelixAdmin; //导入依赖的package包/类
@Override
public List<String> getAssignedInstances(HelixAdmin helixAdmin, String helixClusterName,
    SegmentMetadata segmentMetadata, int numReplicas, String tenantName) {
  String serverTenantName = null;
  if ("realtime".equalsIgnoreCase(segmentMetadata.getIndexType())) {
    serverTenantName = ControllerTenantNameBuilder.getRealtimeTenantNameForTenant(tenantName);
  } else {
    serverTenantName = ControllerTenantNameBuilder.getOfflineTenantNameForTenant(tenantName);
  }
  final Random random = new Random(System.currentTimeMillis());

  List<String> allInstanceList = helixAdmin.getInstancesInClusterWithTag(helixClusterName, serverTenantName);
  List<String> selectedInstanceList = new ArrayList<String>();
  for (int i = 0; i < numReplicas; ++i) {
    final int idx = random.nextInt(allInstanceList.size());
    selectedInstanceList.add(allInstanceList.get(idx));
    allInstanceList.remove(idx);
  }
  LOGGER.info("Segment assignment result for : " + segmentMetadata.getName() + ", in resource : "
      + segmentMetadata.getTableName() + ", selected instances: "
      + Arrays.toString(selectedInstanceList.toArray()));

  return selectedInstanceList;
}
 

import org.apache.helix.HelixAdmin; //导入依赖的package包/类
public static IdealState buildInitialRealtimeIdealStateFor(String realtimeTableName,
    AbstractTableConfig realtimeTableConfig, HelixAdmin helixAdmin, String helixClusterName,
    ZkHelixPropertyStore<ZNRecord> zkHelixPropertyStore) {
  KafkaStreamMetadata kafkaStreamMetadata =
      new KafkaStreamMetadata(realtimeTableConfig.getIndexingConfig().getStreamConfigs());
  String realtimeServerTenant =
      ControllerTenantNameBuilder.getRealtimeTenantNameForTenant(realtimeTableConfig.getTenantConfig().getServer());
  switch (kafkaStreamMetadata.getConsumerType()) {
    case highLevel:
      IdealState idealState =
          buildInitialKafkaHighLevelConsumerRealtimeIdealStateFor(realtimeTableName, helixAdmin, helixClusterName,
              zkHelixPropertyStore);
      List<String> realtimeInstances = helixAdmin.getInstancesInClusterWithTag(helixClusterName, realtimeServerTenant);
      if (realtimeInstances.size() % Integer.parseInt(realtimeTableConfig.getValidationConfig().getReplication()) != 0) {
        throw new RuntimeException("Number of instance in current tenant should be an integer multiples of the number of replications");
      }
      setupInstanceConfigForKafkaHighLevelConsumer(realtimeTableName, realtimeInstances.size(),
          Integer.parseInt(realtimeTableConfig.getValidationConfig().getReplication()), realtimeTableConfig
              .getIndexingConfig().getStreamConfigs(), zkHelixPropertyStore, realtimeInstances);
      return idealState;
    case simple:
    default:
      throw new UnsupportedOperationException("Not support kafka consumer type: "
          + kafkaStreamMetadata.getConsumerType());
  }
}
 

import org.apache.helix.HelixAdmin; //导入依赖的package包/类
public static void main(String[] args) throws Exception {
  if (args.length != 3) {
    System.out.println("USAGE: MySQLAgentLauncher zkAddress clusterName instanceName");
  }
  String zkAddress = args[0];
  String clusterName = args[1];
  String instanceName = args[2];
  HelixAdmin admin = new ZKHelixAdmin(zkAddress);
  List<String> clusters = admin.getClusters();
  System.out.println(clusters);
  if (!clusters.contains(clusterName)) {
    System.err.println("Cluster is not setup. Use cluster-admin to create a  new cluster");
    System.exit(1);
  }
  InstanceConfig instanceConfig = admin.getInstanceConfig(clusterName, instanceName);
  if (instanceConfig == null) {
    String msg =
        "Agent is not yet configured in the cluster. Use cluster-admin to first configure the new Mysql Agent";
    System.err.println(msg);
    System.exit(1);
  }
  admin.close();
  MySQLAgent agent = new MySQLAgent(zkAddress, clusterName, instanceConfig);
  agent.start();
  Thread.currentThread().join();
}
 

import org.apache.helix.HelixAdmin; //导入依赖的package包/类
private void monitor() throws Exception {
  MySQLAdmin mysqlAdmin = _context.getMysqlAdmin();
  HelixAdmin helixAdmin = _context.getHelixManager().getClusterManagmentTool();
  boolean isEnabled = _instanceConfig.getInstanceEnabled();
  int failedPingCount = 0;
  while (true) {
    Thread.sleep(1000);
    if (!mysqlAdmin.ping()) {
      LOG.info("Mysql connection check failed");

      failedPingCount = failedPingCount + 1;
    } else {
      failedPingCount = 0;
      if (!isEnabled) {
        LOG.info("Mysql connection check passed.Enabling the instance");
        helixAdmin.enableInstance(_clusterName, _instanceName, true);
        isEnabled = true;
      }
    }
    if (failedPingCount == 30) {
      LOG.info("Mysql connection check failed 30 times, disabling the instance");
      helixAdmin.enableInstance(_clusterName, _instanceName, false);
      isEnabled = false;
    }
  }
}
 

import org.apache.helix.HelixAdmin; //导入依赖的package包/类
public static void main(String[] args) throws Exception {
  ConsoleAppender ca = new ConsoleAppender();
  ca.setThreshold(Level.INFO);
  ca.setWriter(new OutputStreamWriter(System.out));
  ca.setLayout(new PatternLayout("%-4r [%t] %-5p %c %x - %m%n"));
  Logger.getRootLogger().addAppender(ca);
  String clusterName = "mysql-cluster-test";
  String instanceName = "kgopalak-ld.linkedin.biz_5555";
  String zkAddress = "localhost:2181";
  // TODO:stand alone way to start the agent
  HelixAdmin admin = new ZKHelixAdmin(zkAddress);
  InstanceConfig instanceConfig = admin.getInstanceConfig(clusterName, instanceName);
  MySQLAgent agent = new MySQLAgent(zkAddress, clusterName, instanceConfig);
  agent.start();
  Thread.currentThread().join();
}
 

import org.apache.helix.HelixAdmin; //导入依赖的package包/类
@Test
public void testDisableResource() {
  String className = TestHelper.getTestClassName();
  String methodName = TestHelper.getTestMethodName();
  String clusterName = className + "_" + methodName;
  System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));
  HelixAdmin admin = new ZKHelixAdmin(_gZkClient);
  admin.addCluster(clusterName, true);
  Assert.assertTrue(ZKUtil.isClusterSetup(clusterName, _gZkClient), "Cluster should be setup");
  String resourceName = "TestDB";
  admin.addStateModelDef(clusterName, "MasterSlave", new StateModelDefinition(
      StateModelConfigGenerator.generateConfigForMasterSlave()));
  admin.addResource(clusterName, resourceName, 4, "MasterSlave");
  admin.enableResource(clusterName, resourceName, false);
  BaseDataAccessor<ZNRecord> baseAccessor = new ZkBaseDataAccessor<ZNRecord>(_gZkClient);
  HelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, baseAccessor);
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  IdealState idealState = accessor.getProperty(keyBuilder.idealStates(resourceName));
  Assert.assertFalse(idealState.isEnabled());
  admin.enableResource(clusterName, resourceName, true);
  idealState = accessor.getProperty(keyBuilder.idealStates(resourceName));
  Assert.assertTrue(idealState.isEnabled());
  System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
}
 

import org.apache.helix.HelixAdmin; //导入依赖的package包/类
private static void printStatus(HelixAdmin admin, String cluster, String resource) {
  ExternalView externalView = admin.getResourceExternalView(cluster, resource);
  // System.out.println(externalView);
  TreeSet<String> treeSet = new TreeSet<String>(externalView.getPartitionSet());
  System.out.println("lockName" + "\t" + "acquired By");
  System.out.println("======================================");
  for (String lockName : treeSet) {
    Map<String, String> stateMap = externalView.getStateMap(lockName);
    String acquiredBy = null;
    if (stateMap != null) {
      for (String instanceName : stateMap.keySet()) {
        if ("ONLINE".equals(stateMap.get(instanceName))) {
          acquiredBy = instanceName;
          break;
        }
      }
    }
    System.out.println(lockName + "\t" + ((acquiredBy != null) ? acquiredBy : "NONE"));
  }
}
 

import org.apache.helix.HelixAdmin; //导入依赖的package包/类
@PUT
@Path("{instanceName}")
public Response addInstance(@PathParam("clusterId") String clusterId,
    @PathParam("instanceName") String instanceName, String content) {
  HelixAdmin admin = getHelixAdmin();
  ZNRecord record;
  try {
    record = toZNRecord(content);
  } catch (IOException e) {
    _logger.error("Failed to deserialize user's input " + content + ", Exception: " + e);
    return badRequest("Input is not a vaild ZNRecord!");
  }

  try {
    admin.addInstance(clusterId, new InstanceConfig(record));
  } catch (Exception ex) {
    _logger.error("Error in adding an instance: " + instanceName, ex);
    return serverError(ex);
  }

  return OK();
}
 

import org.apache.helix.HelixAdmin; //导入依赖的package包/类
@PUT
@Path("{instanceName}/configs")
public Response updateInstanceConfig(@PathParam("clusterId") String clusterId,
    @PathParam("instanceName") String instanceName, String content) throws IOException {
  HelixAdmin admin = getHelixAdmin();
  ZNRecord record;
  try {
    record = toZNRecord(content);
  } catch (IOException e) {
    _logger.error("Failed to deserialize user's input " + content + ", Exception: " + e);
    return badRequest("Input is not a vaild ZNRecord!");
  }

  try {
    admin.setInstanceConfig(clusterId, instanceName, new InstanceConfig(record));
  } catch (Exception ex) {
    _logger.error("Error in update instance config: " + instanceName, ex);
    return serverError(ex);
  }

  return OK();
}
 

import org.apache.helix.HelixAdmin; //导入依赖的package包/类
@Override
public List<String> getAssignedInstances(HelixAdmin helixAdmin, ZkHelixPropertyStore<ZNRecord> propertyStore,
    String helixClusterName, SegmentMetadata segmentMetadata, int numReplicas, String tenantName) {
  String serverTenantName = ControllerTenantNameBuilder.getOfflineTenantNameForTenant(tenantName);

  List<String> allInstances = HelixHelper.getEnabledInstancesWithTag(helixAdmin, helixClusterName, serverTenantName);
  List<String> selectedInstanceList = new ArrayList<>();
  if (segmentMetadata.getShardingKey() != null) {
    for (String instance : allInstances) {
      if (HelixHelper.getInstanceConfigsMapFor(instance, helixClusterName, helixAdmin)
          .get("shardingKey")
          .equalsIgnoreCase(segmentMetadata.getShardingKey())) {
        selectedInstanceList.add(instance);
      }
    }
    LOGGER.info("Segment assignment result for : " + segmentMetadata.getName() + ", in resource : "
        + segmentMetadata.getTableName() + ", selected instances: " + Arrays.toString(
        selectedInstanceList.toArray()));
    return selectedInstanceList;
  } else {
    throw new RuntimeException("Segment missing sharding key!");
  }
}
 

import org.apache.helix.HelixAdmin; //导入依赖的package包/类
@Override
public List<String> getAssignedInstances(HelixAdmin helixAdmin, ZkHelixPropertyStore<ZNRecord> propertyStore,
    String helixClusterName, SegmentMetadata segmentMetadata, int numReplicas, String tenantName) {
  String serverTenantName = ControllerTenantNameBuilder.getOfflineTenantNameForTenant(tenantName);
  final Random random = new Random(System.currentTimeMillis());

  List<String> allInstanceList =
      HelixHelper.getEnabledInstancesWithTag(helixAdmin, helixClusterName, serverTenantName);
  List<String> selectedInstanceList = new ArrayList<>();
  for (int i = 0; i < numReplicas; ++i) {
    final int idx = random.nextInt(allInstanceList.size());
    selectedInstanceList.add(allInstanceList.get(idx));
    allInstanceList.remove(idx);
  }
  LOGGER.info("Segment assignment result for : " + segmentMetadata.getName() + ", in resource : "
      + segmentMetadata.getTableName() + ", selected instances: " + Arrays.toString(selectedInstanceList.toArray()));

  return selectedInstanceList;
}
 

import org.apache.helix.HelixAdmin; //导入依赖的package包/类
protected PinotLLCRealtimeSegmentManager(HelixAdmin helixAdmin, String clusterName, HelixManager helixManager,
    ZkHelixPropertyStore propertyStore, PinotHelixResourceManager helixResourceManager, ControllerConf controllerConf,
    ControllerMetrics controllerMetrics) {
  _helixAdmin = helixAdmin;
  _helixManager = helixManager;
  _propertyStore = propertyStore;
  _helixResourceManager = helixResourceManager;
  _clusterName = clusterName;
  _controllerConf = controllerConf;
  _controllerMetrics = controllerMetrics;
  _idealstateUpdateLocks = new Lock[NUM_LOCKS];
  for (int i = 0; i < NUM_LOCKS; i++) {
    _idealstateUpdateLocks[i] = new ReentrantLock();
  }
  _tableConfigCache = new TableConfigCache(_propertyStore);
}
 

import org.apache.helix.HelixAdmin; //导入依赖的package包/类
public static IdealState buildInitialHighLevelRealtimeIdealStateFor(String realtimeTableName,
    TableConfig realtimeTableConfig, HelixAdmin helixAdmin, String helixClusterName,
    ZkHelixPropertyStore<ZNRecord> zkHelixPropertyStore) {
  String realtimeServerTenant =
      ControllerTenantNameBuilder.getRealtimeTenantNameForTenant(realtimeTableConfig.getTenantConfig().getServer());
  final List<String> realtimeInstances = helixAdmin.getInstancesInClusterWithTag(helixClusterName,
      realtimeServerTenant);
  IdealState idealState = buildEmptyKafkaConsumerRealtimeIdealStateFor(realtimeTableName, 1);
  if (realtimeInstances.size() % Integer.parseInt(realtimeTableConfig.getValidationConfig().getReplication()) != 0) {
    throw new RuntimeException(
        "Number of instance in current tenant should be an integer multiples of the number of replications");
  }
  setupInstanceConfigForKafkaHighLevelConsumer(realtimeTableName, realtimeInstances.size(),
      Integer.parseInt(realtimeTableConfig.getValidationConfig().getReplication()),
      realtimeTableConfig.getIndexingConfig().getStreamConfigs(), zkHelixPropertyStore, realtimeInstances);
  return idealState;
}
 

import org.apache.helix.HelixAdmin; //导入依赖的package包/类
HelixAdmin makeHelixAdmin() {
  HelixAdmin admin = mock(HelixAdmin.class);
  ExternalView ev = mock(ExternalView.class);
  IdealState is = mock(IdealState.class);
  when(admin.getResourceExternalView(clusterName, tableName)).thenReturn(ev);
  when(admin.getResourceIdealState(clusterName, tableName)).thenReturn(is);

  List<String> segmentsInIs = segmentsInIdealStateOrExtView();
  Map<String, String> dummy = new HashMap<>(1);
  dummy.put("someHost", "ONLINE");
  for (String segment : segmentsInIs) {
    when(is.getInstanceStateMap(segment)).thenReturn(dummy);
  }
  when(ev.getStateMap(anyString())).thenReturn(null);

  return admin;
}
 

import org.apache.helix.HelixAdmin; //导入依赖的package包/类
public void testBulkDeleteWithFailures(boolean useSet) throws Exception {
  HelixAdmin helixAdmin =  makeHelixAdmin();
  ZkHelixPropertyStore<ZNRecord> propertyStore = makePropertyStore();
  FakeDeletionManager deletionManager = new FakeDeletionManager(helixAdmin, propertyStore);
  Collection<String> segments;
  if (useSet) {
    segments = new HashSet<String>();
  } else {
    segments = new ArrayList<String>();
  }

  segments.addAll(segmentsThatShouldBeDeleted());
  segments.addAll(segmentsInIdealStateOrExtView());
  segments.addAll(segmentsFailingPropStore());
  deletionManager.deleteSegmentsFromPropertyStoreAndLocal(tableName, segments);

  Assert.assertTrue(deletionManager.segmentsToRetry.containsAll(segmentsFailingPropStore()));
  Assert.assertTrue(deletionManager.segmentsToRetry.containsAll(segmentsInIdealStateOrExtView()));
  Assert.assertTrue(deletionManager.segmentsRemovedFromStore.containsAll(segmentsThatShouldBeDeleted()));
}