org.apache.helix.tools.ClusterSetup Java Examples

private static void setup() {

    IDefaultNameSpace defaultNameSpace = new IDefaultNameSpace() {
      @Override
      public void createDefaultNameSpace(ZkClient client) {
        client.deleteRecursive("/" + clusterName);

      }
    };

    zkServer = new ZkServer(dataDir, logDir, defaultNameSpace, port);
    zkServer.start();
    ClusterSetup clusterSetup = new ClusterSetup(zkConnectString);
    clusterSetup.setupTestCluster(clusterName);
    try {
      Thread.sleep(1000);
    } catch (InterruptedException e) {
      e.printStackTrace();
    }
  }
 

/**
 * Test cluster creation according to the pre-set routing mapping.
 * Helix Java API tested is ClusterSetup in this method.
 */
@Test
public void testCreateClusters() {
  // Create two ClusterSetups using two different constructors
  // Note: ZK Address here could be anything because multiZk mode is on (it will be ignored)
  ClusterSetup clusterSetupZkAddr = new ClusterSetup(ZK_SERVER_MAP.keySet().iterator().next());
  ClusterSetup clusterSetupBuilder = new ClusterSetup.Builder().build();

  createClusters(clusterSetupZkAddr);
  verifyClusterCreation(clusterSetupZkAddr);

  createClusters(clusterSetupBuilder);
  verifyClusterCreation(clusterSetupBuilder);

  // Create clusters again to continue with testing
  createClusters(clusterSetupBuilder);
}
 

@Test()
public void testDisableNode() throws Exception {
  String command =
      "-zkSvr " + ZK_ADDR + " -enableInstance " + CLUSTER_NAME + " " + PARTICIPANT_PREFIX
          + "_12918" + " TestDB TestDB_0 false";
  ClusterSetup.processCommandLineArgs(command.split(" "));
  boolean result =
      ClusterStateVerifier.verifyByPolling(new ClusterStateVerifier.BestPossAndExtViewZkVerifier(
          ZK_ADDR, CLUSTER_NAME));
  Assert.assertTrue(result);

  ZKHelixAdmin tool = new ZKHelixAdmin(_gZkClient);
  tool.enableInstance(CLUSTER_NAME, PARTICIPANT_PREFIX + "_12918", true);

  result =
      ClusterStateVerifier.verifyByPolling(new ClusterStateVerifier.BestPossAndExtViewZkVerifier(
          ZK_ADDR, CLUSTER_NAME));
  Assert.assertTrue(result);

}
 

@Override
public Representation delete() {
  try {
    String clusterName =
        ResourceUtil.getAttributeFromRequest(getRequest(), ResourceUtil.RequestKey.CLUSTER_NAME);
    String instanceName =
        ResourceUtil.getAttributeFromRequest(getRequest(), ResourceUtil.RequestKey.INSTANCE_NAME);
    ZkClient zkclient =
        ResourceUtil.getAttributeFromCtx(getContext(), ResourceUtil.ContextKey.ZKCLIENT);
    ClusterSetup setupTool = new ClusterSetup(zkclient);
    setupTool.dropInstanceFromCluster(clusterName, instanceName);
    getResponse().setStatus(Status.SUCCESS_OK);
  } catch (Exception e) {
    getResponse().setEntity(ClusterRepresentationUtil.getErrorAsJsonStringFromException(e),
        MediaType.APPLICATION_JSON);
    getResponse().setStatus(Status.SUCCESS_OK);
    LOG.error("Error in delete instance", e);
  }
  return null;
}
 

StringRepresentation getClustersRepresentation() throws JsonGenerationException,
    JsonMappingException, IOException {
  ZkClient zkClient =
      ResourceUtil.getAttributeFromCtx(getContext(), ResourceUtil.ContextKey.ZKCLIENT);

  ClusterSetup setupTool = new ClusterSetup(zkClient);
  List<String> clusters = setupTool.getClusterManagementTool().getClusters();

  ZNRecord clustersRecord = new ZNRecord("Clusters Summary");
  clustersRecord.setListField("clusters", clusters);
  StringRepresentation representation =
      new StringRepresentation(ClusterRepresentationUtil.ZNRecordToJson(clustersRecord),
          MediaType.APPLICATION_JSON);

  return representation;
}
 

private void verifyClusterSetupMsdsEndpoint(
    RealmAwareZkClient.RealmAwareZkConnectionConfig connectionConfig) {
  System.out.println("Start " + TestHelper.getTestMethodName());

  ClusterSetup firstClusterSetup = new ClusterSetup.Builder().build();
  ClusterSetup secondClusterSetup =
      new ClusterSetup.Builder().setRealmAwareZkConnectionConfig(connectionConfig).build();

  try {
    verifyMsdsZkRealm(CLUSTER_ONE, true,
        () -> firstClusterSetup.addCluster(CLUSTER_ONE, false));
    verifyMsdsZkRealm(CLUSTER_FOUR, false,
        () -> firstClusterSetup.addCluster(CLUSTER_FOUR, false));

    verifyMsdsZkRealm(CLUSTER_FOUR, true,
        () -> secondClusterSetup.addCluster(CLUSTER_FOUR, false));
    verifyMsdsZkRealm(CLUSTER_ONE, false,
        () -> secondClusterSetup.addCluster(CLUSTER_ONE, false));
  } finally {
    firstClusterSetup.close();
    secondClusterSetup.close();
  }
}
 

@Override
public Representation delete() {
  try {
    String clusterName =
        ResourceUtil.getAttributeFromRequest(getRequest(), ResourceUtil.RequestKey.CLUSTER_NAME);
    String resourceName =
        ResourceUtil.getAttributeFromRequest(getRequest(), ResourceUtil.RequestKey.RESOURCE_NAME);
    ZkClient zkclient =
        ResourceUtil.getAttributeFromCtx(getContext(), ResourceUtil.ContextKey.ZKCLIENT);

    ClusterSetup setupTool = new ClusterSetup(zkclient);
    setupTool.dropResourceFromCluster(clusterName, resourceName);
    getResponse().setStatus(Status.SUCCESS_OK);
  } catch (Exception e) {
    getResponse().setEntity(ClusterRepresentationUtil.getErrorAsJsonStringFromException(e),
        MediaType.APPLICATION_JSON);
    getResponse().setStatus(Status.SUCCESS_OK);
    LOG.error("", e);
  }
  return null;
}
 

@BeforeClass
public void beforeClass() throws Exception {
  ZkServer zkServer = TestHelper.startZkServer(_zkAddr);
  _zkServerRef.set(zkServer);
  _zkClient = SharedZkClientFactory.getInstance()
      .buildZkClient(new HelixZkClient.ZkConnectionConfig(_zkAddr));
  _zkClient.setZkSerializer(new ZNRecordSerializer());
  _setupTool = new ClusterSetup(_zkClient);
  _participants = new MockParticipantManager[_numNodes];
  _setupTool.addCluster(CLUSTER_NAME, true);
  setupParticipants(_setupTool);
  setupDBs(_setupTool);
  createManagers(_zkAddr, CLUSTER_NAME);

  // start controller
  String controllerName = CONTROLLER_PREFIX + "_0";
  _controller = new ClusterControllerManager(_zkAddr, CLUSTER_NAME, controllerName);
  _controller.syncStart();

  ZkHelixClusterVerifier clusterVerifier =
      new BestPossibleExternalViewVerifier.Builder(CLUSTER_NAME).setZkAddr(_zkAddr).build();
  Assert.assertTrue(clusterVerifier.verifyByPolling());
}
 

@BeforeSuite
public void beforeSuite() throws Exception {
  // TODO: use logging.properties file to config java.util.logging.Logger levels
  java.util.logging.Logger topJavaLogger = java.util.logging.Logger.getLogger("");
  topJavaLogger.setLevel(Level.WARNING);

  // start zk
  _zkServer = TestHelper.startZkServer(ZK_ADDR);
  AssertJUnit.assertTrue(_zkServer != null);
  ZKClientPool.reset();

  _gZkClient =
      new ZkClient(ZK_ADDR, ZkClient.DEFAULT_SESSION_TIMEOUT, ZkClient.DEFAULT_CONNECTION_TIMEOUT,
          new ZNRecordSerializer());
  _gSetupTool = new ClusterSetup(_gZkClient);

  // start admin
  _adminThread = new AdminThread(ZK_ADDR, ADMIN_PORT);
  _adminThread.start();

  // create a client
  _gClient = new Client(Protocol.HTTP);

  // wait for the web service to start
  Thread.sleep(100);
}
 

StringRepresentation getSchedulerTasksRepresentation() throws JsonGenerationException,
    JsonMappingException, IOException {
  String clusterName = (String) getRequest().getAttributes().get("clusterName");
  String instanceName = (String) getRequest().getAttributes().get("instanceName");
  ZkClient zkClient = (ZkClient) getContext().getAttributes().get(RestAdminApplication.ZKCLIENT);
  ClusterSetup setupTool = new ClusterSetup(zkClient);
  List<String> instances =
      setupTool.getClusterManagementTool().getInstancesInCluster(clusterName);

  HelixDataAccessor accessor =
      ClusterRepresentationUtil.getClusterDataAccessor(zkClient, clusterName);
  LiveInstance liveInstance =
      accessor.getProperty(accessor.keyBuilder().liveInstance(instanceName));
  String sessionId = liveInstance.getEphemeralOwner();

  StringRepresentation representation = new StringRepresentation("");// (ClusterRepresentationUtil.ObjectToJson(instanceConfigs),
                                                                     // MediaType.APPLICATION_JSON);

  return representation;
}
 

@Test()
public void testDropResource() throws Exception {
  // add a resource to be dropped
  _gSetupTool.addResourceToCluster(CLUSTER_NAME, "MyDB", 6, STATE_MODEL);
  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, "MyDB", 3);

  boolean result =
      ClusterStateVerifier.verifyByPolling(new ClusterStateVerifier.BestPossAndExtViewZkVerifier(
          ZK_ADDR, CLUSTER_NAME));
  Assert.assertTrue(result);

  String command = "-zkSvr " + ZK_ADDR + " -dropResource " + CLUSTER_NAME + " " + "MyDB";
  ClusterSetup.processCommandLineArgs(command.split(" "));

  TestHelper.verifyWithTimeout("verifyEmptyCurStateAndExtView", 30 * 1000, CLUSTER_NAME, "MyDB",
      TestHelper.<String> setOf("localhost_12918", "localhost_12919", "localhost_12920",
          "localhost_12921", "localhost_12922"), ZK_ADDR);
}
 

/**
 * Remove constraints
 * <p>
 * Usage:
 * <code>curl -X DELETE http://{host:port}/clusters/{cluster}/constraints/MESSAGE_CONSTRAINT/{constraintId}
 */
@Override
public Representation delete() {
  String clusterName =
      ResourceUtil.getAttributeFromRequest(getRequest(), ResourceUtil.RequestKey.CLUSTER_NAME);
  String constraintTypeStr =
      ResourceUtil.getAttributeFromRequest(getRequest(), ResourceUtil.RequestKey.CONSTRAINT_TYPE);
  String constraintId =
      ResourceUtil.getAttributeFromRequest(getRequest(), ResourceUtil.RequestKey.CONSTRAINT_ID);

  try {
    ZkClient zkClient =
        ResourceUtil.getAttributeFromCtx(getContext(), ResourceUtil.ContextKey.ZKCLIENT);
    ClusterSetup setupTool = new ClusterSetup(zkClient);

    setupTool.removeConstraint(clusterName, constraintTypeStr, constraintId);
  } catch (Exception e) {
    LOG.error("Error in delete constraint", e);
    getResponse().setEntity(ClusterRepresentationUtil.getErrorAsJsonStringFromException(e),
        MediaType.APPLICATION_JSON);
    getResponse().setStatus(Status.SUCCESS_OK);
  }
  return null;
}
 

/**
 * Remove a cluster
 * <p>
 * Usage: <code> curl -X DELETE http://{host:port}/clusters/{clusterName}
 */
@Override
public Representation delete() {
  try {
    String clusterName =
        ResourceUtil.getAttributeFromRequest(getRequest(), ResourceUtil.RequestKey.CLUSTER_NAME);
    ZkClient zkClient =
        ResourceUtil.getAttributeFromCtx(getContext(), ResourceUtil.ContextKey.ZKCLIENT);
    ClusterSetup setupTool = new ClusterSetup(zkClient);
    setupTool.deleteCluster(clusterName);
    getResponse().setStatus(Status.SUCCESS_OK);
  } catch (Exception e) {
    getResponse().setEntity(ClusterRepresentationUtil.getErrorAsJsonStringFromException(e),
        MediaType.APPLICATION_JSON);
    getResponse().setStatus(Status.SUCCESS_OK);
  }
  return null;
}
 

private static void addConfiguration(ClusterSetup setup, String baseDir, String clusterName,
    String instanceName) throws IOException {
  Map<String, String> properties = new HashMap<String, String>();
  HelixConfigScopeBuilder builder = new HelixConfigScopeBuilder(ConfigScopeProperty.PARTICIPANT);
  HelixConfigScope instanceScope =
      builder.forCluster(clusterName).forParticipant(instanceName).build();
  properties.put("change_log_dir", baseDir + instanceName + "/translog");
  properties.put("file_store_dir", baseDir + instanceName + "/filestore");
  properties.put("check_point_dir", baseDir + instanceName + "/checkpoint");
  setup.getClusterManagementTool().setConfig(instanceScope, properties);
  FileUtils.deleteDirectory(new File(properties.get("change_log_dir")));
  FileUtils.deleteDirectory(new File(properties.get("file_store_dir")));
  FileUtils.deleteDirectory(new File(properties.get("check_point_dir")));
  new File(properties.get("change_log_dir")).mkdirs();
  new File(properties.get("file_store_dir")).mkdirs();
  new File(properties.get("check_point_dir")).mkdirs();
}
 

@Override
@BeforeClass
public void beforeClass() throws Exception {
  System.out.println("START " + CLASS_NAME + " at " + new Date(System.currentTimeMillis()));

  // setup storage cluster
  _gSetupTool.addCluster(CLUSTER_NAME, true);
  _gSetupTool.addResourceToCluster(CLUSTER_NAME, TEST_DB, _PARTITIONS, STATE_MODEL);

  for (int i = 0; i < NODE_NR; i++) {
    String storageNodeName = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    _gSetupTool.addInstanceToCluster(CLUSTER_NAME, storageNodeName);
  }
  _gSetupTool.rebalanceStorageCluster(CLUSTER_NAME, TEST_DB, 3);

  // Set the timeout values
  IdealState idealState =
      _gSetupTool.getClusterManagementTool().getResourceIdealState(CLUSTER_NAME, TEST_DB);
  String stateTransition = "SLAVE" + "-" + "MASTER" + "_" + Message.Attributes.TIMEOUT;
  idealState.getRecord().setSimpleField(stateTransition, "300");

  String command =
      "-zkSvr " + ZK_ADDR + " -addResourceProperty " + CLUSTER_NAME + " " + TEST_DB + " "
          + stateTransition + " 200";
  ClusterSetup.processCommandLineArgs(command.split(" "));
}
 

@AfterSuite
public void afterSuite() throws IOException {
  // Clean up all JMX objects
  for (ObjectName mbean : _server.queryNames(null, null)) {
    try {
      _server.unregisterMBean(mbean);
    } catch (Exception e) {
      // OK
    }
  }

  // Close all ZK resources
  _baseDataAccessorMap.values().forEach(BaseDataAccessor::close);
  _clusterSetupMap.values().forEach(ClusterSetup::close);
  _helixZkClientMap.values().forEach(HelixZkClient::close);
  _zkServerMap.values().forEach(TestHelper::stopZkServer);
}
 

/**
 * Starts an additional in-memory ZooKeeper for testing.
 * @param i index to be added to the ZK port to avoid conflicts
 * @throws Exception
 */
private void startZooKeeper(int i)
    throws Exception {
  String zkAddress = ZK_PREFIX + (ZK_START_PORT + i);
  ZkServer zkServer = TestHelper.startZkServer(zkAddress);
  AssertJUnit.assertNotNull(zkServer);
  HelixZkClient.ZkClientConfig clientConfig = new HelixZkClient.ZkClientConfig();
  clientConfig.setZkSerializer(new ZNRecordSerializer());
  HelixZkClient zkClient = DedicatedZkClientFactory.getInstance()
      .buildZkClient(new HelixZkClient.ZkConnectionConfig(zkAddress), clientConfig);
  ClusterSetup gSetupTool = new ClusterSetup(zkClient);
  BaseDataAccessor baseDataAccessor = new ZkBaseDataAccessor<>(zkClient);

  _zkServerMap.put(zkAddress, zkServer);
  _helixZkClientMap.put(zkAddress, zkClient);
  _clusterSetupMap.put(zkAddress, gSetupTool);
  _baseDataAccessorMap.put(zkAddress, baseDataAccessor);
}
 

public static void dropCluster(String clusterName, RealmAwareZkClient zkClient, ClusterSetup setup) {
  String namespace = "/" + clusterName;
  if (zkClient.exists(namespace)) {
    try {
      setup.deleteCluster(clusterName);
    } catch (Exception ex) {
      // Failed to delete, give some more time for connections to drop
      try {
        Thread.sleep(3000L);
        setup.deleteCluster(clusterName);
      } catch (Exception ignored) {
        // OK - just ignore
      }
    }
  }
}
 

/**
 * Create a Helix cluster for the Gobblin Cluster application.
 *
 * @param zkConnectionString the ZooKeeper connection string
 * @param clusterName the Helix cluster name
 * @param overwrite true to overwrite exiting cluster, false to reuse existing cluster
 */
public static void createGobblinHelixCluster(String zkConnectionString, String clusterName, boolean overwrite) {
  ClusterSetup clusterSetup = new ClusterSetup(zkConnectionString);
  // Create the cluster and overwrite if it already exists
  clusterSetup.addCluster(clusterName, overwrite);
  // Helix 0.6.x requires a configuration property to have the form key=value.
  String autoJoinConfig = ZKHelixManager.ALLOW_PARTICIPANT_AUTO_JOIN + "=true";
  clusterSetup.setConfig(HelixConfigScope.ConfigScopeProperty.CLUSTER, clusterName, autoJoinConfig);
}
 

@Override
public Representation post(Representation entity) {
  try {
    String clusterName =
        ResourceUtil.getAttributeFromRequest(getRequest(), ResourceUtil.RequestKey.CLUSTER_NAME);
    String resourceName =
        ResourceUtil.getAttributeFromRequest(getRequest(), ResourceUtil.RequestKey.RESOURCE_NAME);
    ZkClient zkclient =
        ResourceUtil.getAttributeFromCtx(getContext(), ResourceUtil.ContextKey.ZKCLIENT);
    ClusterSetup setupTool = new ClusterSetup(zkclient);

    JsonParameters jsonParameters = new JsonParameters(entity);
    String command = jsonParameters.getCommand();
    if (command.equalsIgnoreCase(ClusterSetup.resetResource)) {

      setupTool.getClusterManagementTool()
          .resetResource(clusterName, Arrays.asList(resourceName));
    } else if (command.equalsIgnoreCase(ClusterSetup.enableResource)) {
      jsonParameters.verifyCommand(ClusterSetup.enableResource);
      boolean enabled = Boolean.parseBoolean(jsonParameters.getParameter(JsonParameters.ENABLED));
      setupTool.getClusterManagementTool().enableResource(clusterName, resourceName, enabled);
    } else {
      throw new HelixException("Unsupported command: " + command + ". Should be one of ["
          + ClusterSetup.resetResource + "]");
    }
  } catch (Exception e) {
    getResponse().setEntity(ClusterRepresentationUtil.getErrorAsJsonStringFromException(e),
        MediaType.APPLICATION_JSON);
    getResponse().setStatus(Status.SUCCESS_OK);
    LOG.error("", e);
  }
  return null;
}
 

private Map<String, String> addInstanceCmd(String instances) {
  Map<String, String> parameters = new HashMap<String, String>();
  parameters.put(JsonParameters.INSTANCE_NAMES, instances);
  parameters.put(JsonParameters.MANAGEMENT_COMMAND, ClusterSetup.addInstance);

  return parameters;
}
 

/**
 * Enable/disable Helix controller
 * <p>
 * Usage:
 * <code>curl -d 'jsonParameters={"command":"enableCluster","enabled":"{true/false}"}'
 * -H "Content-Type: application/json" http://{host:port}/clusters/{cluster}/Controller
 */
@Override
public Representation post(Representation entity) {
  try {
    String clusterName =
        ResourceUtil.getAttributeFromRequest(getRequest(), ResourceUtil.RequestKey.CLUSTER_NAME);
    ZkClient zkClient =
        ResourceUtil.getAttributeFromCtx(getContext(), ResourceUtil.ContextKey.ZKCLIENT);
    ClusterSetup setupTool = new ClusterSetup(zkClient);

    JsonParameters jsonParameters = new JsonParameters(entity);
    String command = jsonParameters.getCommand();

    if (command == null) {
      throw new HelixException("Could NOT find 'command' in parameterMap: "
          + jsonParameters._parameterMap);
    } else if (command.equalsIgnoreCase(ClusterSetup.enableCluster)) {
      boolean enabled = Boolean.parseBoolean(jsonParameters.getParameter(JsonParameters.ENABLED));

      setupTool.getClusterManagementTool().enableCluster(clusterName, enabled);
    } else {
      throw new HelixException("Unsupported command: " + command + ". Should be one of ["
          + ClusterSetup.enableCluster + "]");
    }

    getResponse().setEntity(getControllerRepresentation(clusterName));
    getResponse().setStatus(Status.SUCCESS_OK);

  } catch (Exception e) {
    getResponse().setEntity(ClusterRepresentationUtil.getErrorAsJsonStringFromException(e),
        MediaType.APPLICATION_JSON);
    getResponse().setStatus(Status.SUCCESS_OK);
  }

  return null;
}
 

/**
 * Add a new Helix cluster
 * <p>
 * Usage: <code> curl -d 'jsonParameters={"command":"addCluster","clusterName":"{clusterName}"}' -H
 * "Content-Type: application/json" http://{host:port}/clusters
 */
@Override
public Representation post(Representation entity) {
  try {
    JsonParameters jsonParameters = new JsonParameters(entity);
    String command = jsonParameters.getCommand();

    if (command == null) {
      throw new HelixException("Could NOT find 'command' in parameterMap: "
          + jsonParameters._parameterMap);
    } else if (command.equalsIgnoreCase(ClusterSetup.addCluster)) {
      jsonParameters.verifyCommand(ClusterSetup.addCluster);

      ZkClient zkClient =
          ResourceUtil.getAttributeFromCtx(getContext(), ResourceUtil.ContextKey.ZKCLIENT);
      ClusterSetup setupTool = new ClusterSetup(zkClient);
      setupTool.addCluster(jsonParameters.getParameter(JsonParameters.CLUSTER_NAME), false);
    } else {
      throw new HelixException("Unsupported command: " + command + ". Should be one of ["
          + ClusterSetup.addCluster + "]");
    }

    getResponse().setEntity(getClustersRepresentation());
    getResponse().setStatus(Status.SUCCESS_OK);
  } catch (Exception e) {
    getResponse().setEntity(ClusterRepresentationUtil.getErrorAsJsonStringFromException(e),
        MediaType.APPLICATION_JSON);
    getResponse().setStatus(Status.SUCCESS_OK);
    LOG.error("Error in posting " + entity, e);
  }
  return null;
}
 

@Test
public void test() throws Exception {
  String jsonPayload =
      "{\"command\":\"resetPartition\",\"resource\": \"DB-1\",\"partition\":\"DB-1_22 DB-1_23\"}";
  Map<String, String> map = ClusterRepresentationUtil.JsonToMap(jsonPayload);
  Assert.assertNotNull(map.get(JsonParameters.MANAGEMENT_COMMAND));
  Assert.assertEquals(ClusterSetup.resetPartition, map.get(JsonParameters.MANAGEMENT_COMMAND));
  Assert.assertNotNull(map.get(JsonParameters.RESOURCE));
  Assert.assertEquals("DB-1", map.get(JsonParameters.RESOURCE));
  Assert.assertNotNull(map.get(JsonParameters.PARTITION));
  Assert.assertEquals("DB-1_22 DB-1_23", map.get(JsonParameters.PARTITION));
}
 

void assertClusterSetupException(String command) {
  boolean exceptionThrown = false;
  try {
    ClusterSetup.processCommandLineArgs(command.split(" "));
  } catch (Exception e) {
    exceptionThrown = true;
  }
  Assert.assertTrue(exceptionThrown);
}
 

public static void setupCluster(String clusterName, String zkAddr, int startPort,
    String participantNamePrefix, String resourceNamePrefix, int resourceNb, int partitionNb,
    int nodesNb, int replica, String stateModelDef, RebalanceMode mode, boolean doRebalance)
    throws Exception {
  HelixZkClient zkClient = SharedZkClientFactory.getInstance()
      .buildZkClient(new HelixZkClient.ZkConnectionConfig(zkAddr));
  if (zkClient.exists("/" + clusterName)) {
    LOG.warn("Cluster already exists:" + clusterName + ". Deleting it");
    zkClient.deleteRecursively("/" + clusterName);
  }

  ClusterSetup setupTool = new ClusterSetup(zkAddr);
  setupTool.addCluster(clusterName, true);

  for (int i = 0; i < nodesNb; i++) {
    int port = startPort + i;
    setupTool.addInstanceToCluster(clusterName, participantNamePrefix + "_" + port);
  }

  for (int i = 0; i < resourceNb; i++) {
    String resourceName = resourceNamePrefix + i;
    setupTool.addResourceToCluster(clusterName, resourceName, partitionNb, stateModelDef,
        mode.toString());
    if (doRebalance) {
      setupTool.rebalanceStorageCluster(clusterName, resourceName, replica);
    }
  }
  zkClient.close();
}
 

private Map<String, String> enablePartitionCmd(String resourceName, String partitions,
    boolean enabled) {
  Map<String, String> parameters = new HashMap<String, String>();
  parameters.put(JsonParameters.MANAGEMENT_COMMAND, ClusterSetup.enablePartition);
  parameters.put(JsonParameters.ENABLED, "" + enabled);
  parameters.put(JsonParameters.PARTITION, partitions);
  parameters.put(JsonParameters.RESOURCE, resourceName);

  return parameters;
}
 

@BeforeClass
public void beforeClass() throws Exception {
  System.out.println("START " + CLASS_NAME + " at " + new Date(System.currentTimeMillis()));

  _setupTool = new ClusterSetup(ZK_ADDR);
  // setup storage cluster
  _setupTool.addCluster(CLUSTER_NAME, true);
  _setupTool.addResourceToCluster(CLUSTER_NAME, TEST_DB, NUM_PARTITIONS, STATE_MODEL);

  for (int i = 0; i < NUM_PARTICIPANTS; i++) {
    String storageNodeName = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    _setupTool.addInstanceToCluster(CLUSTER_NAME, storageNodeName);
  }
  _setupTool.rebalanceStorageCluster(CLUSTER_NAME, TEST_DB, NUM_REPLICAS);

  // start dummy participants
  for (int i = 0; i < NUM_PARTICIPANTS; i++) {
    String instanceName = PARTICIPANT_PREFIX + "_" + (START_PORT + i);
    _participants[i] = new MockParticipantManager(ZK_ADDR, CLUSTER_NAME, instanceName);
    _participants[i].syncStart();
  }

  // start controller
  String controllerName = CONTROLLER_PREFIX + "_0";
  _controller = new ClusterControllerManager(ZK_ADDR, CLUSTER_NAME, controllerName);
  _controller.syncStart();

  boolean result = ClusterStateVerifier.verifyByPolling(
      new ClusterStateVerifier.MasterNbInExtViewVerifier(ZK_ADDR, CLUSTER_NAME), 10000, 100);
  Assert.assertTrue(result);

  result = ClusterStateVerifier.verifyByPolling(
      new ClusterStateVerifier.BestPossAndExtViewZkVerifier(ZK_ADDR, CLUSTER_NAME), 10000, 100);
  Assert.assertTrue(result);

  // create cluster manager
  _manager = HelixManagerFactory.getZKHelixManager(CLUSTER_NAME, "Admin",
      InstanceType.ADMINISTRATOR, ZK_ADDR);
  _manager.connect();
}
 

@Override
public Representation post(Representation entity) {
  try {
    String clusterName = (String) getRequest().getAttributes().get("clusterName");
    ZkClient zkClient =
        (ZkClient) getContext().getAttributes().get(RestAdminApplication.ZKCLIENT);
    ;

    JsonParameters jsonParameters = new JsonParameters(entity);
    String command = jsonParameters.getCommand();

    if (command.equalsIgnoreCase(ClusterSetup.addStateModelDef)) {
      ZNRecord newStateModel =
          jsonParameters.getExtraParameter(JsonParameters.NEW_STATE_MODEL_DEF);
      HelixDataAccessor accessor =
          ClusterRepresentationUtil.getClusterDataAccessor(zkClient, clusterName);

      accessor.setProperty(accessor.keyBuilder().stateModelDef(newStateModel.getId()),
          new StateModelDefinition(newStateModel));
      getResponse().setEntity(getStateModelsRepresentation());
    } else {
      throw new HelixException("Unsupported command: " + command + ". Should be one of ["
          + ClusterSetup.addStateModelDef + "]");
    }

    getResponse().setStatus(Status.SUCCESS_OK);
  } catch (Exception e) {
    getResponse().setEntity(ClusterRepresentationUtil.getErrorAsJsonStringFromException(e),
        MediaType.APPLICATION_JSON);
    getResponse().setStatus(Status.SUCCESS_OK);
    LOG.error("Error in posting " + entity, e);
  }
  return null;
}
 

private Map<String, String> rebalanceCmd(int replicas, String prefix, String tag) {
  Map<String, String> parameters = new HashMap<String, String>();
  parameters.put(JsonParameters.REPLICAS, "" + replicas);
  if (prefix != null) {
    parameters.put(JsonParameters.RESOURCE_KEY_PREFIX, prefix);
  }
  if (tag != null) {
    parameters.put(ClusterSetup.instanceGroupTag, tag);
  }
  parameters.put(JsonParameters.MANAGEMENT_COMMAND, ClusterSetup.rebalance);

  return parameters;
}